flax

2f2ec1a0 - Return sharding whenever the leaf has .sharding attribute, not checking explicitly for jax.Array.