Copies a tf.Tensor onto the DTensor device with the given layout.
tf.experimental.dtensor.copy_to_mesh(
tensor: Any,
layout: tf.experimental.dtensor.Layout
,
source_layout: Optional[tf.experimental.dtensor.Layout
] = None
) -> tf.Tensor
Copies a regular tf.Tensor onto the DTensor device. Use the mesh attached to
layout
as target mesh. This method currently only supports replicated
layouts. To get a DTensor with a sharded layout, use the pack
method.
Args | |
---|---|
tensor
|
A regular tf.Tensor to be copied as a DTensor. |
layout
|
Target layout (and mesh) for the result DTensor. |
source_layout
|
Source layout of the tensor before copy, used for backward passes. |
Returns | |
---|---|
A DTensor on the DTensor device with the given layout. |