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 | 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