jax
342d62dd - Merge pull request #4848 from apaszke:pmap-in-axes

Commit

5 years ago

Merge pull request #4848 from apaszke:pmap-in-axes PiperOrigin-RevId: 341624889

References

#4848 - Add support for pmap in_axes other than 0 and None

Author

a-googler

a-googler

Parents

Loading