jax

e2cc5689 - raise type errors consistently in jaxpr checker