jax

768a64eb - Cache `canonicalize_sharding` and remove the canonicalization in `_scatter_update` since the out_sharding passed to scatter_update should already be