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, если макет представляет собой устройство без распределения.

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