PR #5020 Use jax.shard_map in flax

Use jax.shard_map in flax #5020

copybara-service merged 1 commit into main from test_817877943

copybara-service force pushed from 666684f1 to 12027e47 172 days ago

Use jax.shard_map in flax

5bf9b35f

copybara-service force pushed from 12027e47 to 5bf9b35f 172 days ago

copybara-service merged 5bf9b35f into main 172 days ago

copybara-service deleted the test_817877943 branch 172 days ago

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone