tree 92ab1b7a74fbcad4f895e3e2e8632dabef3f5541
parent 9908d143206dffc4b24b9b028ebfa0dda392e05d
author A. Unique TensorFlower <gardener@tensorflow.org> 1597087198 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> 1597087810 -0700

Implement AllPermute collective in TF2

AllPermute takes
- a list of devices participating in the collective
- a permutation as a list of integers.
- a tensor

The list of devices replaces the need for group_key and group_size. The number of inputs only scales with the number of devices within one group.

The integers in the permutation are based on indices of the list of devices.
E.g. devices = {"GPU:0", "GPU:1"} and permutation = {1,0} means
- devices[0] sends to devices[permutation[0]] and
- devices[1] sends to devices[permutation[1]].

PiperOrigin-RevId: 325860722
Change-Id: I919438fa1cd75d684e47f927289de389e8dcb3f0