jax.make_mesh#
- jax.make_mesh(axis_shapes, axis_names, *, devices=None, axis_types=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 ) 会影响设备的排序。
如果您想使用它提供的额外参数（如 contiguous_submeshes 和 allow_split_physical_axes），可以使用 jax.experimental.mesh_utils.create_device_mesh 。