Changes the layout of tensor to the same as layout_tensor.
tf.experimental.dtensor.relayout_like(
tensor: tf.Tensor,
layout_tensor: tf.Tensor,
name: Optional[str] = None
) -> tf.Tensor
relayout_like is often used inside a tf.function, to ensure a tensor is placed to the same mesh and with the same layout as another tensor.
The backward gradient of a relayout is a relayout_like operation, to ensure the backward tensor has the same layout as the forward input tensor:
@ops.RegisterGradient("Relayout")
def _relayout_gradient(op, grad):
return relayout_like(grad, layout_input=op.inputs[0])
Here is another illustrative example:
@tf.function def func(x): z = tf.ones(x.shape) z = dtensor.relayout_like(z, x) return x + z with dtensor.default_mesh(cpu_mesh): x = tf.ones((4, 4)) with dtensor.default_mesh(gpu_mesh): y = func(x) # y would be on the cpu mesh, following the mesh of x.
| Args | |
|---|---|
tensor | A DTensor to specify a new layout for. |
layout_tensor | A Tensor object whose layout will be used for the layout of result. The shape and type of layout_tensor are irrelevant. |
name | name of the Op. |
| Returns | |
|---|---|
| A DTensor output from the RelayoutLike op. |
© 2022 The TensorFlow Authors. All rights reserved.
Licensed under the Creative Commons Attribution License 4.0.
Code samples licensed under the Apache 2.0 License.
https://www.tensorflow.org/api_docs/python/tf/experimental/dtensor/relayout_like