jax
43549157 - Merge pull request #4904 from apaszke:pmap-out-axes

Commit

5 years ago

Merge pull request #4904 from apaszke:pmap-out-axes PiperOrigin-RevId: 344104708

References

#4904 - Add support for out_axes in pmap

Author

a-googler

a-googler

Parents

Loading