flax
[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.
#5374

Merged

[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`. #5374

copybara-service merged 1 commit into main from test_888582438

copybara-service force pushed from 640ecfe6 to 7124c65e 4 days ago

copybara-service force pushed from 7124c65e to 41d00520 4 days ago

copybara-service force pushed from 41d00520 to 4ddd1a36 4 days ago

[pmap] In-line definitions of `jax.device_put_sharded` and `jax.devic…

355e141a

copybara-service force pushed from 4ddd1a36 to 355e141a 4 days ago

copybara-service merged 355e141a into main 4 days ago

copybara-service deleted the test_888582438 branch 4 days ago

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone