tf.raw_ops.CollectivePermute

An Op to permute tensors across replicated TPU instances.

tf.raw_ops.CollectivePermute(
    input, source_target_pairs, name=None
)

Each instance supplies its own input.

For example, suppose there are 4 TPU instances: [A, B, C, D]. Passing source_target_pairs=[[0,1],[1,2],[2,3],[3,0]] gets the outputs: [D, A, B, C].

input A Tensor. Must be one of the following types: float32, float64, int32, uint8, int16, int8, complex64, int64, qint8, quint8, qint32, bfloat16, uint16, complex128, half, uint32, uint64. The local input to be permuted. Currently only supports float and bfloat16.
source_target_pairs A Tensor of type int32. A tensor with shape [num_pairs, 2].
name A name for the operation (optional).

A Tensor. Has the same type as input.