tf.raw_ops.CollectivePermute
Stay organized with collections Save and categorize content based on your preferences.
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 , qint16 , quint16 , 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 . |