def testForiLoopErrors(self):
    'Test typing error messages for while.'
    with self.assertRaisesRegex(TypeError, 'arguments to fori_loop must have equal types'):
        lax.fori_loop(np.int16(0), jnp.int32(10), (lambda i, c: c), jnp.float32(7))
Test typing error messages for while.
def testForiLoopErrors(self):
    
    with self.assertRaisesRegex(TypeError, 'arguments to fori_loop must have equal types'):
        lax.fori_loop(np.int16(0), jnp.int32(10), (lambda i, c: c), jnp.float32(7))