jax

44c13598 - Remove the redundant context manager in lower_fun. We already enter into eqn's context one level above in jaxpr_subcomp