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.