jax

69d7bcf7 - except-and-raise during jaxpr checking, adding jaxpr as context, and simplify type environment