tf.raw_ops.CollectivePermute
Операция для перестановки тензоров между реплицированными экземплярами TPU.
tf.raw_ops.CollectivePermute(
input, source_target_pairs, name=None
)
Каждый экземпляр предоставляет свой собственный вход.
Например, предположим, что есть 4 экземпляра TPU: [A, B, C, D]. Передача source_target_pairs=[[0,1],[1,2],[2,3],[3,0]] приводит к следующим выходным значениям: [D, A, B, C].
| Аргументы | |
|---|---|
input | A Tensor. Должен быть одним из следующих типов: float32, float64, int32, uint8, int16, int8, complex64, int64, qint8, quint8, qint32, bfloat16, uint16, complex128, half, uint32, uint64. Местный вход для перестановки. В настоящее время поддерживаются только float и bfloat16. |
source_target_pairs | A Tensor типа int32. Тензор с формой [num_pairs, 2]. |
name | Имя операции (необязательно). |
| Возвращает | |
|---|---|
A Tensor. Имеет тот же тип, что и input. |
© 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/versions/r2.9/api_docs/python/tf/raw_ops/CollectivePermute