PR #3016 Upgrade the `pjit` guide to `jax.jit`

Upgrade the `pjit` guide to `jax.jit` #3016

copybara-service merged 2 commits into google:main from IvyZX:jit-guide

IvyZX requested a review from

cgarciae 2 years ago

8bitmp3 requested a review from

8bitmp3 2 years ago

IvyZX force pushed 2 years ago

upgrade jit guide

acce3a28

Update Scale up Flax Modules on multiple devices

87e7fa75

IvyZX force pushed to 87e7fa75 2 years ago

8bitmp3 assigned

8bitmp3 2 years ago

8bitmp3 added pull ready

copybara-service merged bdb68ef5 into main 2 years ago

Reviewers

cgarciae

8bitmp3

Assignees

8bitmp3

Labels

pull ready

Milestone

No milestone