jax

cc34ed26 - check aval compatibility, not strict equality, when typechecking jaxpr equations