jax

9ffdd6bd - Add a type check that verifies the lower and upper arguments to lax.fori_loop have equal types. (#1693)