tf.experimental.dtensor.relayout_like | TensorFlow v2.16.1 (original) (raw)

tf.experimental.dtensor.relayout_like

Stay organized with collections Save and categorize content based on your preferences.

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.

Except as otherwise noted, the content of this page is licensed under the Creative Commons Attribution 4.0 License, and code samples are licensed under the Apache 2.0 License. For details, see the Google Developers Site Policies. Java is a registered trademark of Oracle and/or its affiliates. Some content is licensed under the numpy license.

Last updated 2024-04-26 UTC.