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

返回要使用的编译选项，这些选项源于标志值。

参数：

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

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

  • device_assignment – 可选的 jax 设备数组，指示逻辑副本到物理设备的分配（默认继承自 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

目录