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])