flax
4e2becd8
- Preserve sharding information in axes_scan
Go
Login via GitHub
Home
Pricing
FAQ
Install
Login
via GitHub
Commit
View On
GitHub
Commit
239 days ago
Preserve sharding information in axes_scan Also, remove jnp.array call in normalization.py, to preserve sharding information PiperOrigin-RevId: 780954856
References
to-linen-eval-shape
#4806 - Preserve sharding information in axes_scan
Author
oliverdutton-iso
Committer
a-googler
Parents
cbb4f80f
Loading