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 на сетке unpack возвращает сырые компоненты, размещенные на каждом базовом устройстве.
2. Упаковки этих сырых компонентов в том же порядке с использованием pack возвращает DTensor, который должен быть идентичен исходному DTensor — как по значению содержимого, так и по расположению.

Форма, ранг и скаляры: ранг DTensor совпадает с рангом его сырых компонентов, т. е. ранг сохраняется. Это приводит к согласованной интерпретации для упаковки скалярных значений в DTensor. Единственное допустимое расположение для скалярного значения — полностью реплицированное, и отдельные компоненты должны быть идентичными скалярами.

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

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

Например, предположим, что у нас есть сетка [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])