jax

b532e068 - Add recursive _prefetch_jaxpr_literals method, use it in both xla and pxla.