tf.experimental.dtensor.pack

Упаковывает компоненты tf.Tensor в DTensor.

tf.experimental.dtensor.pack(
    tensors: Sequence[Any],
    layout: tf.experimental.dtensor.Layout
) -> Any

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

Упаковка и распаковка — обратные операции:

* unpack(pack(tensors)) == tensors
* pack(unpack(dtensor)) == dtensor

1. Для любого DTensor на mesh функция unpack возвращает исходные компоненты, размещённые на каждом базовом устройстве.
2. Упаковка этих исходных компонентов в том же порядке с помощью pack возвращает DTensor, который должен быть идентичен исходному DTensor — как по значению содержимого, так и по компоновке.

Форма, ранг и скалярные значения: Ранг DTensor совпадает с рангом его исходных компонентов, то есть ранг сохраняется. Благодаря этому скалярные значения в DTensor интерпретируются согласованно. Для скалярного значения допустима только полностью реплицированная компоновка, а отдельные компоненты должны быть идентичными скалярами.

Каждый входной tensors[i] будет скопирован на layout.mesh.local_device[i], если он ещё не находится на локальном устройстве. Не следует передавать нелокальные компоненты в pack; используйте copy_to_mesh и relayout, чтобы разместить тензоры на всех глобальных устройствах mesh.

Вызывающий код отвечает за то, чтобы базовые значения для pack соответствовали указанной компоновке и чтобы число заданных значений не превышало число локальных устройств. Pack не перемещает данные между клиентами. Дополнительные сведения о компоновках см. в примерах ниже.

Например, предположим, что у нас есть mesh [X(2), Y(3)], состоящий в общей сложности из 6 базовых устройств. Кроме того, предположим, что сопоставление расположений устройств выглядит следующим образом:

device_ID  |  location X, Y
        0     0, 0
        1     0, 1
        2     0, 2
        3     1, 0
        4     1, 1
        5     1, 2

1. Для одномерного векторного DTensor с формой [128], компоновкой [mesh.X] и значением range(128) исходные компоненты будут иметь форму [64] каждый, а сами исходные компоненты будут выглядеть так:

   device_ID  |  raw component
           0     range(0, 64)
           1     range(0, 64)
           2     range(0, 64)
           3     range(64, 128)
           4     range(64, 128)
           5     range(64, 128)
   

   Это также означает, что для одномерного DTensor с формой [2] и компоновкой [mesh.X] исходные компоненты имеют форму [1], а не форму скалярных значений [].

2. Для двумерного векторного DTensor с формой [2, 3], компоновкой [mesh.X, mesh.Y] и значением range(6) получается полностью сегментированный DTensor.

   В глобальном представлении содержимое выглядит так:

   [
     [0.0, 1.0, 2.0],
     [3.0, 4.0, 5.0],
   ]
   

   Исходные компоненты будут иметь форму [1, 1] каждый и следующее содержимое:

   device_ID  |  raw component
           0     [[0.0]]
           1     [[1.0]]
           2     [[2.0]]
           3     [[3.0]]
           4     [[4.0]]
           5     [[5.0]]
   

3. Для скалярного значения 123.0 DTensor допустима только одна компоновка [] (без измерений, но с полной репликацией).

   Исходные компоненты будут иметь форму [] каждый и следующее содержимое:

   device_ID  |  raw component
           0     123.0
           1     123.0
           2     123.0
           3     123.0
           4     123.0
           5     123.0
   

   И снова вызывающий код pack должен предоставить 6 идентичных исходных компонентов со скалярной формой.

4. Для трёхмерного векторного DTensor с формой [2, 2, 3], компоновкой [X, unsharded, unsharded] и значением range(12)

   В глобальном представлении содержимое выглядит так:

   [
     [
       [0.0, 1.0, 2.0],
       [3.0, 4.0, 5.0],
     ],
     [
       [6.0, 7.0, 8.0],
       [9.0, 10., 11.],
     ],
   ]
   

   Исходные компоненты будут иметь форму [1, 2, 3] каждый и следующее содержимое:

   device_ID  |  raw component
           0     range(6).reshape([1, 2, 3])
           1     range(6).reshape([1, 2, 3])
           2     range(6).reshape([1, 2, 3])
           3     range(6, 12).reshape([1, 2, 3])
           4     range(6, 12).reshape([1, 2, 3])
           5     range(6, 12).reshape([1, 2, 3])