all_gather for TPU doesn't support backward gradients.

[tchaton]

Currently, we rely on AllGatherGrad to compute gather for GPUs.

TODO:

• [] Extend this class to support TPU
• [] Add tests

[ethanwharris]

Although all_gather will now work, it still doesn't support grad. so re-opening :)