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

copybara-service[bot] merged 3 commits intogoogle:maingoogle/flax:mainfrom

samanklesaria:eval_shape_shardingsamanklesaria/flax:eval_shape_shardingCopy head branch name to clipboard