jax

a9e9f97f - Use no_tracing config in `_create_pjit_jaxpr` to so that AOT path can also error if we re-trace.