jax
a7e5eaee - [pallas:mosaic_gpu] `jnp.reduce_sum` now works for >1D arrays

Commit

289 days ago

[pallas:mosaic_gpu] `jnp.reduce_sum` now works for >1D arrays PiperOrigin-RevId: 737578598

References

Author

superbobry

Committer

Google-ML-Automation

Parents