flax

2e1f5adc - Merge pull request #5205 from samanklesaria:embed_out_sharding