flax

Added sharding propagation support in nnx.eval_shape

#5111

Closed