@parameterized.named_parameters(({'testcase_name': f'_dtype={dtype.__name__}', 'dtype': dtype} for dtype in jtu.dtypes.all_integer))
def test_while_loop_init_weak_type(self, dtype):

    def cond_fun(val):
        return (val < 2)

    def body_fun(val):
        return (val + increment)
    increment = jnp.array(1, dtype=dtype)
    init_weak = 0
    result = lax.while_loop(cond_fun, body_fun, init_weak)
    self.assertArraysEqual(result, jnp.full_like(increment, 2))
@parameterized.named_parameters(({'testcase_name': f'_dtype={dtype.__name__}', 'dtype': dtype} for dtype in jtu.dtypes.all_integer))
def test_while_loop_init_weak_type(self, dtype):

    def cond_fun(val):
        return (val < 2)

    def body_fun(val):
        return (val + increment)
    increment = jnp.array(1, dtype=dtype)
    init_weak = 0
    result = lax.while_loop(cond_fun, body_fun, init_weak)
    self.assertArraysEqual(result, jnp.full_like(increment, 2))