jax.experimental.mesh_utils 模組

用於建構裝置網格的工具。

API

create_device_mesh(mesh_shape[, devices, ...])	為 jax.sharding.Mesh 建立高效能的裝置網格。
create_hybrid_device_mesh(mesh_shape, ...[, ...])	為混合式（例如，ICI 和 DCN）平行處理建立裝置網格。