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].
Args | |
|---|---|
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). |
Returns | |
|---|---|
A Tensor. Has the same type as input.
|