jax.lib.xla_bridge.get_compile_options

jax.lib.xla_bridge.get_compile_options#

jax.lib.xla_bridge.get_compile_options(num_replicas, num_partitions, device_assignment=None, use_spmd_partitioning=True, use_shardy_partitioner=False, use_auto_spmd_partitioning=False, auto_spmd_partitioning_mesh_shape=None, auto_spmd_partitioning_mesh_ids=None, env_options_overrides=None, fdo_profile=None, detailed_logging=True, backend=None)[源代码][源代码]#

返回根据标志值推导出的编译选项。

参数:

  • num_replicas (int) – 要编译的副本数量。

  • num_partitions (int) – 要编译的分区数量。

  • device_assignment – 可选的 jax 设备 ndarray，指示逻辑副本到物理设备的分配（默认继承自 xla_client.CompileOptions）。必须与 num_replicas 和 num_partitions 一致。

  • use_spmd_partitioning (bool) – 布尔值，指示是否在XLA中启用SPMD或MPMD分区。

  • use_shardy_partitioner (bool) – 布尔值，指示是否在 XLA 中使用 Shardy 分区器。Shardy 是一个新的开源 MLIR 传播框架。目前 Shardy 在 JAX 中是实验性的。请参阅 www.github.com/openxla/shardy。

  • use_auto_spmd_partitioning (bool) – 布尔值，指示是否为SPMD分区器自动生成XLA分片。

  • auto_spmd_partitioning_mesh_shape (list[int] | None) – 用于创建 auto_spmd_partitioning 搜索空间的设备网格形状。

  • auto_spmd_partitioning_mesh_ids (list[int] | None) – 用于创建 auto_spmd_partitioning 搜索空间的设备 ID。

  • env_options_overrides (dict[str, str] | None) – 由编译器解析的附加选项的字典

  • fdo_profile (bytes | None) – 传递给XLA的反馈导向优化的可选配置文件。

  • detailed_logging (bool) – 这是一个关于 XLA 是否应该明智地记录编译信息的“有趣”计算吗？

  • backend (xc.Client | None) – 客户端，如果可用的话。

返回类型:

xc.CompileOptions

目录