jax

dcd0eb59 - Replace pe.trace_to_jaxpr_dynamic with pe.trace_to_jaxpr in ad_checkpoint.py