jax
fa38f250 - Add support for all_to_all over vmapped axes

Commit

5 years ago

Add support for all_to_all over vmapped axes

References

#4429 - Add support for `all_to_all` in vmap

Author

apaszke

apaszke

Committer

apaszke

apaszke

Parents

Loading