tf.experimental.dtensor.relayout

Изменяет раскладку tensor.

tf.experimental.dtensor.relayout(
    tensor: tf.Tensor,
    layout: tf.experimental.dtensor.Layout,
    name: Optional[str] = None
) -> tf.Tensor

Используется в блокнотах

Изменяет раскладку tensor на layout. Это используется для тонкой настройки поведения операций, следующих за tensor или связанных с ним, например для выбора одного шаблона расширения SPMD вместо другого. Это работает за счет прямого распространения layout на связанные графы вычислений TensorFlow при распространении раскладки.

В настоящее время поддерживается только преобразование раскладок из реплицированной в шардированную или из шардированной в реплицированную по каждому измерению сетки. То есть поддерживается преобразование "x, y" -> "unsharded, y", но не поддерживается "x, y" -> "z, y".

Также поддерживается специальная спецификация шардинга "match", которая указывает, что relayout должен выполнять тождественную операцию для любого шардинга по этим измерениям сетки.

Внутри relayout преобразуется в набор операций Split и/или AllToAll. При преобразовании раскладок тензора из реплицированной в шардированную затраты относительно невелики, поскольку добавляются только операции Split и межустройственное взаимодействие не требуется. Однако при преобразовании раскладок тензора из шардированной в реплицированную может происходить межустройственное взаимодействие, что потенциально снижает производительность.