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

Loading