flax
8fe8f5eb - Plumb spmd_axis_name from vmap_with_axes through to JAX vmap

Commit

3 years ago

Plumb spmd_axis_name from vmap_with_axes through to JAX vmap PiperOrigin-RevId: 467276339

References

Author

James Lee-Thorp

Committer

a-googler

Parents