flax

Add sharding propagation support in nnx.eval_shape (clone of #5111)

#5247

Merged