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