jax
1bf80368 - Merge pull request #4170 from apaszke:all_to_all

Commit

5 years ago

Merge pull request #4170 from apaszke:all_to_all PiperOrigin-RevId: 333502317

References

#4170 - Allow vmapping all_to_all and implement a (slow) CPU and GPU translation

Author

a-googler

a-googler

Parents

Loading