tf.experimental.dtensor.copy_to_mesh

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.