jax.make_mesh#

jax.make_mesh(axis_shapes, axis_names, axis_types=None, *, devices=None)[源]#

根据指定的形状和轴名称创建高效网格。

此函数尝试自动计算从一组逻辑轴到物理网格的良好映射。例如，在具有 8 个设备的 TPU v3 上

>>> mesh = jax.make_mesh((8,), ('x'))  
>>> [d.id for d in mesh.devices.flat]  
[0, 1, 2, 3, 6, 7, 4, 5]

上述排序考虑了 TPU v3 的物理拓扑。它将设备排列成一个环形，这在 TPU v3 上产生了高效的全规约操作。

现在，让我们看看另一个具有 16 个 TPU v3 设备的示例

>>> mesh = jax.make_mesh((2, 8), ('x', 'y'))  
>>> [d.id for d in mesh.devices.flat]  
[0, 1, 2, 3, 6, 7, 4, 5, 8, 9, 10, 11, 14, 15, 12, 13]
>>> mesh = jax.make_mesh((4, 4), ('x', 'y'))  
>>> [d.id for d in mesh.devices.flat]  
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]

如您所见，逻辑轴 (axis_shapes) 会影响设备的排序。

如果您想使用 jax.experimental.mesh_utils.create_device_mesh 提供的额外参数，例如 contiguous_submeshes 和 allow_split_physical_axes，则可以使用它。

参数:

  • axis_shapes (Sequence[int]) – 网格的形状。例如，axis_shape=(4, 2)

  • axis_names (Sequence[str]) – 网格轴的名称。例如，axis_names=(‘x’, ‘y’)

  • axis_types (tuple[mesh_lib.AxisType, ...] | None) – 可选的 jax.sharding.AxisType 条目元组，对应于 axis_names。有关更多信息，请参阅 显式分片。

  • devices (Sequence[xc.Device] | None) – 可选的仅关键字参数，允许您指定要创建网格的设备。

返回:

一个 jax.sharding.Mesh 对象。

返回类型:

mesh_lib.Mesh