tf.experimental.dtensor.Layout

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

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

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

Например, рассмотрим сетку 1-D:

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

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

Для другого примера рассмотрим сетку 2-D:

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() -> layout_pb2.LayoutProto

Создаёт представление прототипа макета.

batch_sharded

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

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

Возвращает макет, фрагментированный по размерности партии.

delete

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

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

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

from_str

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

@staticmethod
from_str(
    layout_str: bytes
) -> 'Layout'

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

from_string

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

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

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

inner_sharded

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

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

Возвращает макет, фрагментированный по внутренней размерности.

is_fully_replicated

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

is_fully_replicated() -> bool

Возвращает True, если все оси тензора дублируются.

mesh_proto

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

mesh_proto() -> layout_pb2.MeshProto

Возвращает базовую сетку в формате протокола.

num_shards

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

num_shards(
    idx: int
) -> int

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

offset_to_shard

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

offset_to_shard()

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

offset_tuple_to_global_index

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

offset_tuple_to_global_index(
    offset_tuple
)

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

replicated

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

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

Возвращает дублированный макет ранга rank.

serialized_string

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

serialized_string() -> bytes

Возвращает сериализованное двоичное строковое представление протокола.

to_string

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

to_string() -> str

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

unravel

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

unravel(
    unpacked_tensors: List[np.ndarray]
) -> np.ndarray

Преобразует сплющенный список фрагментов в фрагментированный массив.

__eq__

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

__eq__(
    other
) -> bool

Возвращает self==value.