[NVIDIA TF] Add DTensorAllToAll and relayout implementation
@bfontain @rainwoodman @pjannaty  FYI: @RuijiaoSun  Thanks for the PR. The changes looked good to me. 

I'll pull it internally to give the XLA part of the logic a try on TPUs. Good news! The tests for the XLA code path passed on our V3 TPU machines. 

This PR adds a new high-level op, `DTensorAllToAll` along with lowering to  `CollectiveAllToAllV2`.
To start, I've implemented a relayout using this new all-to-all for certain cases only. There should be many more applications for all-to-all that we can implement in the future.

## Supported relayout cases
I've added support for the cases where one axis is becoming sharded to a certain mesh dimension while another axis is becoming unsharded from that mesh dimension. For example, `[x,unsharded] -> [unsharded, x]` or `[unsharded, y, x] -> [x, y, unsharded]`. You can look at the unit tests in `spmd_test.py` for more (non-exhaustive) examples.

Before this PR, these relayouts would be implemented using all-gather to fully replicate the tensor, followed by slice. With all-to-all, the communication overhead is reduced since only the minimal slices of data that need to be moved are exchanged across the devices.

## Summary of changes
* DTensorAllToAll op
* DTensorAllToAllLowering pass which can lower to `CollectiveAllToAllV2` for cpu and gpu, or XLA AllToAll for TPU. **I did not test the TPU path at all.**
* Use DTensorAllToAll for certain relayouts
* MLIR test
* Python unit tests