tf.experimental.dtensor.relayout_like

Изменяет компоновку tensor так, чтобы она соответствовала компоновке layout_tensor.

tf.experimental.dtensor.relayout_like(
    tensor: tf.Tensor,
    layout_tensor: tf.Tensor,
    name: Optional[str] = None
) -> tf.Tensor

relayout_like часто используется внутри tf.function, чтобы гарантировать, что тензор размещен в той же сетке и имеет ту же компоновку, что и другой тензор.

Обратный градиент relayout — это операция relayout_like, которая гарантирует, что обратный тензор имеет ту же компоновку, что и входной тензор прямого прохода:

@ops.RegisterGradient("Relayout")
def _relayout_gradient(op, grad):
  return relayout_like(grad, layout_input=op.inputs[0])

Вот еще один наглядный пример:

@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.

Аргументы
`tensor`	DTensor, для которого требуется указать новую компоновку.
`layout_tensor`	Объект Tensor, компоновка которого будет использоваться для компоновки результата. Форма и тип layout_tensor не имеют значения.
`name`	имя операции.

Возвращает
Выходной DTensor, возвращаемый операцией RelayoutLike.

© 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