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

    def func(carry, x):
        return ((carry + x), x)
    init_weak = 0
    x = jnp.ones(5, dtype=dtype)
    (carry, result) = lax.scan(func, init_weak, x)
    self.assertEqual(carry, x.sum())
    self.assertArraysEqual(result, x)
@parameterized.named_parameters(({'testcase_name': f'_dtype={dtype.__name__}', 'dtype': dtype} for dtype in jtu.dtypes.all_integer))
def test_scan_init_weak_type(self, dtype):

    def func(carry, x):
        return ((carry + x), x)
    init_weak = 0
    x = jnp.ones(5, dtype=dtype)
    (carry, result) = lax.scan(func, init_weak, x)
    self.assertEqual(carry, x.sum())
    self.assertArraysEqual(result, x)