Packs tf.Tensor
components into a DTensor.
tf.experimental.dtensor.pack(
tensors: Sequence[Any],
layout: tf.experimental.dtensor.Layout
) -> Any
Used in the notebooks
Used in the guide | Used in the tutorials |
---|---|
Packing and unpacking are inverse operations:
* unpack(pack(tensors)) == tensors
* pack(unpack(dtensor)) == dtensor
- For any DTensor on the mesh,
unpack
returns the raw components placed on each underlying device. - Packing these raw components in the same order using
pack
returns a DTensor which should be identical to the original DTensor--both the content value and the layout.
Shape, Rank, and Scalars: The rank of the DTensor is the same as the rank of its raw components, i.e., rank is preserved. This leads to a consistent interpretation for packing scalar values into a DTensor. The only valid layout for a scalar value is fully replicated, and the individual components must be identical scalars.
Each input tensors[i]
will be copied to layout.mesh.local_device[i]
if not already on the local device. Non-local components should not be passed
to pack
; use copy_to_mesh
and relayout
to place tensors on all global
devices on a mesh.
It is the caller's responsibility to ensure that the underlying values
for pack
adhere to the specified layout, and that only as many values are
specified as there are local devices. Pack does not move data between clients.
See examples below for more detail about layouts.
For example, assume we have a mesh [X(2), Y(3)]
, which has in total 6
underlying devices. Futuremore, assume that the device location mapping is
the following:
device_ID | location X, Y
0 0, 0
1 0, 1
2 0, 2
3 1, 0
4 1, 1
5 1, 2
For 1-D vector DTensor with shape
[128]
with layout[mesh.X]
and value asrange(128)
, the raw components will have shape[64]
each, and the raw components will be: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)
This also means for a 1-D DTensor with shape
[2]
and layout[mesh.X]
, the raw components have shape[1]
rather than the shape for scalar values[]
.For 2-D vector DTensor with shape
[2, 3]
with layout[mesh.X, mesh.Y]
and value asrange(6)
, this is basically a fully-sharded DTensor.From global view, the content looks like
[ [0.0, 1.0, 2.0], [3.0, 4.0, 5.0], ]
The raw components will have shape
[1, 1]
each, and have the following content:device_ID | raw component 0 [[0.0]] 1 [[1.0]] 2 [[2.0]] 3 [[3.0]] 4 [[4.0]] 5 [[5.0]]
For a scalar value
123.0
DTensor, it can only have one legitimate layout[]
(no dimension, but fully replicated).The raw components will have shape
[]
each, and have the following content:device_ID | raw component 0 123.0 1 123.0 2 123.0 3 123.0 4 123.0 5 123.0
Again, caller of
pack
is expected to provide 6 identical value raw components with scalar shapes.For 3-D vector DTensor with shape
[2, 2, 3]
with layout[X, unsharded, unsharded]
and value asrange(12)
,From global view, the content looks like:
[ [ [0.0, 1.0, 2.0], [3.0, 4.0, 5.0], ], [ [6.0, 7.0, 8.0], [9.0, 10., 11.], ], ]
The raw components will have shape
[1, 2, 3]
each, and have the following content: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])
Args | |
---|---|
tensors
|
The list of local tensor components to pack into a DTensor. |
layout
|
The layout of the DTensor to be created. |
Returns | |
---|---|
A DTensor created from the individual component tensors. |