jax

060bd8a4 - tidy jaxpr typechecking error test