tf.experimental.dtensor.Layout

Представляет информацию о размещении DTensor.

tf.experimental.dtensor.Layout(
    sharding_specs: List[str],
    mesh: tf.experimental.dtensor.Mesh
)

Используется в блокнотах

Размещение описывает, как распределенный тензор разбивается между сеткой (а значит, и между устройствами). Для каждой оси тензора соответствующая спецификация разбиения указывает, по какому измерению сетки выполняется разбиение. Специальная спецификация разбиения UNSHARDED указывает, что эта ось реплицируется на всех устройствах данной сетки.

Подробное обсуждение и примеры приведены в разделе «Основные понятия DTensor».

Например, рассмотрим одномерную сетку:

Mesh(["TPU:0", "TPU:1", "TPU:2", "TPU:3", "TPU:4", "TPU:5"], [("x", 6)])

Эта сетка располагает 6 устройств TPU в одномерном массиве. Layout([UNSHARDED], mesh) — это размещение тензора ранга 1, который реплицируется на 6 устройствах.

Рассмотрим еще один пример — двумерную сетку:

Mesh(["TPU:0", "TPU:1", "TPU:2", "TPU:3", "TPU:4", "TPU:5"],
     [("x", 3), ("y", 2)])

Эта сетка располагает 6 устройств TPU в двумерном массиве 3x2. Layout(["x", UNSHARDED], mesh) — это размещение тензора ранга 2, первая ось которого разбита по измерению сетки «x», а вторая реплицирована. Если разместить np.arange(6).reshape((3, 2)) с помощью этого размещения, отдельные тензоры-компоненты будут выглядеть так:

Device  |  Component
 TPU:0     [[0, 1]]
 TPU:1     [[0, 1]]
 TPU:2     [[2, 3]]
 TPU:3     [[2, 3]]
 TPU:4     [[4, 5]]
 TPU:5     [[4, 5]]

Методы

as_proto

as_proto()

as_proto(self: tensorflow.python._pywrap_dtensor_device.Layout) -> tensorflow::dtensor::LayoutProto

Возвращает сообщение-протобуфер LayoutProto.

batch_sharded

Посмотреть исходный код

@classmethod
batch_sharded(
    mesh: tf.experimental.dtensor.Mesh,
    batch_dim: str,
    rank: int,
    axis: int = 0
) -> 'Layout'

Возвращает размещение с разбиением по измерению пакета.

delete

Посмотреть исходный код

delete(
    dims: List[int]
) -> 'Layout'

Возвращает размещение, из которого удалены указанные измерения.

from_device

Посмотреть исходный код

@classmethod
from_device(
    device: str
) -> 'Layout'

Создает размещение для одного устройства на основе сетки из одного устройства.

from_proto

Посмотреть исходный код

@classmethod
from_proto(
    layout_proto: layout_pb2.LayoutProto
) -> 'Layout'

Создает экземпляр на основе LayoutProto.

from_single_device_mesh

Посмотреть исходный код

@classmethod
from_single_device_mesh(
    mesh: tf.experimental.dtensor.Mesh
) -> 'Layout'

Создает размещение для одного устройства на основе сетки из одного устройства.

from_string

Посмотреть исходный код

@classmethod
from_string(
    layout_str: str
) -> 'Layout'

Создает экземпляр на основе строки, понятной человеку.

global_shape_from_local_shape

global_shape_from_local_shape()

global_shape_from_local_shape(self: tensorflow.python._pywrap_dtensor_device.Layout, local_shape: List[int]) -> tuple

Возвращает глобальную форму, вычисленную на основе этой локальной формы.

inner_sharded

Посмотреть исходный код

@classmethod
inner_sharded(
    mesh: tf.experimental.dtensor.Mesh,
    inner_dim: str,
    rank: int
) -> 'Layout'

Возвращает раскладку, разбитую на шарды по внутреннему измерению.

is_batch_parallel

is_batch_parallel()

is_batch_parallel(self: tensorflow.python._pywrap_dtensor_device.Layout) -> bool

is_fully_replicated

is_fully_replicated()

is_fully_replicated(self: tensorflow.python._pywrap_dtensor_device.Layout) -> bool

Возвращает True, если все оси тензора реплицированы.

is_single_device

is_single_device()

is_single_device(self: tensorflow.python._pywrap_dtensor_device.Layout) -> bool

Возвращает True, если Layout представляет нераспределённое устройство.

local_shape_from_global_shape

local_shape_from_global_shape()

local_shape_from_global_shape(self: tensorflow.python._pywrap_dtensor_device.Layout, global_shape: List[int]) -> tuple

Возвращает локальную форму, вычисленную на основе этой глобальной формы.

num_shards

num_shards()

num_shards(self: tensorflow.python._pywrap_dtensor_device.Layout, idx: int) -> int

Возвращает количество шардов для измерения тензора idx.

offset_to_shard

Посмотреть исходный код

offset_to_shard()

Сопоставление смещения в плоском списке с индексом шарда.

offset_tuple_to_global_index

Посмотреть исходный код

offset_tuple_to_global_index(
    offset_tuple
)

Сопоставление смещения с индексом в глобальном тензоре.

replicated

Посмотреть исходный код

@classmethod
replicated(
    mesh: tf.experimental.dtensor.Mesh,
    rank: int
) -> 'Layout'

Возвращает реплицированную раскладку ранга rank.

to_parted

Посмотреть исходный код

to_parted() -> 'Layout'

Возвращает «разделённую» раскладку на основе статической раскладки.

Разделённая раскладка содержит оси, которые большинство расширителей SPMD рассматривают как независимые.

FIXME(b/285905569): Точная семантика всё ещё исследуется.

to_string

to_string()

to_string(self: tensorflow.python._pywrap_dtensor_device.Layout) -> str

__eq__

__eq__()

eq(self: tensorflow.python._pywrap_dtensor_device.Layout, arg0: tensorflow.python._pywrap_dtensor_device.Layout) -> bool