tf.experimental.dtensor.create_tpu_mesh

Возвращает распределённую сетку TPU, оптимизированную для кольцевых редукций AllReduce.

tf.experimental.dtensor.create_tpu_mesh(
    mesh_dim_names: List[str],
    mesh_shape: List[int],
    mesh_name: str,
    ring_dims: Optional[int] = None,
    ring_axes: Optional[List[str]] = None,
    ring_bounds: Optional[List[int]] = None,
    can_split_host_across_rings: bool = True,
    build_ring_across_rings: bool = False,
    rotate_ring_across_rings: bool = False,
    use_xla_spmd: bool = layout_lib.USE_XLA_SPMD
) -> tf.experimental.dtensor.Mesh

Для построения колец будут использоваться только необходимые ведущие оси, указанные в ring_axes, при условии, что подмассив, образованный этими осями, содержит достаточно ядер для кольца требуемого размера. Оставшиеся оси в ring_axes не повлияют на результаты.

Эта функция всегда использует все устройства TPU и предоставляет больше возможностей настройки, чем tf.experimental.dtensor.create_distributed_mesh.