jax
c0cfc7ae - Merge pull request #21032 from mattjj:vmap-grad-shmap-bug

Commit

1 year ago

Merge pull request #21032 from mattjj:vmap-grad-shmap-bug PiperOrigin-RevId: 630375950

References

Author

a-googler

Parents