tree 3a4e9b2e038c140af868f62c86721ce782e2de14
parent 89df3ddcd5451aefe52c7deb684f220b0520a6b1
author Ran Chen <crccw@google.com> 1592338274 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> 1592338830 -0700

Explicitly take the set of devices in CollectiveAllReduce

We used to infer the devices from the inputs, but sometimes the inputs don't
have device placement. E.g. when passing into or returning from tf.function, the
device placement may be lost.

Instead of inferring from the inputs we should just be explicit about the
collective devices.

PiperOrigin-RevId: 316743112
Change-Id: I2f6995f2f4cc86864723e203deb7562363cdbc38