Added sharding propagation support in nnx.eval_shape#5111

vfdev-5 wants to merge 1 commit intomaingoogle/flax:mainfrom

eval-shape-shardinggoogle/flax:eval-shape-shardingCopy head branch name to clipboard