flax

6bce1438 - Support explicit sharding for nnx.eval_shape