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)[source]#

返回要使用的编译选项，这些选项是从标志值派生的。

参数:

  • 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 | None) – 用于创建 auto_spmd_partitioning 搜索空间的设备网格形状。

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

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

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

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

  • backend (xc.Client | None | None) – 客户端（如果可用）。

返回类型:

xc.CompileOptions