jax

fc18e57b - remove `pjit_p`, leaving only `jax.extend.core.primitives.jit_p`