@parameterized.named_parameters(({'testcase_name': '_jit_scan={}_jit_f={}_impl={}'.format(jit_scan, jit_f, scan_name), 'jit_scan': jit_scan, 'jit_f': jit_f, 'scan': scan_impl} for jit_scan in [False, True] for jit_f in [False, True] for (scan_impl, scan_name) in SCAN_IMPLS))
def testScanLinearize(self, jit_scan, jit_f, scan):
    rng = np.random.RandomState(0)
    d = rng.randn(2)

    def f(c, a):
        assert (a.shape == (3,))
        assert (c.shape == (4,))
        b = jnp.cos(((jnp.sum(jnp.sin(a)) + jnp.sum(jnp.cos(c))) + jnp.sum(jnp.tan(d))))
        c = jnp.sin((c * b))
        assert (b.shape == ())
        return (c, b)
    if jit_f:
        f = api.jit(f)
    if jit_scan:
        scan = api.jit(scan, static_argnums=(0,))
    as_ = rng.randn(5, 3)
    c = rng.randn(4)
    ans = api.linearize((lambda c, as_: scan(f, c, as_)), c, as_)[1](c, as_)
    expected = api.linearize((lambda c, as_: scan_reference(f, c, as_)), c, as_)[1](c, as_)
    self.assertAllClose(ans, expected, check_dtypes=False, rtol={np.float64: 1e-14})
@parameterized.named_parameters(({'testcase_name': '_jit_scan={}_jit_f={}_impl={}'.format(jit_scan, jit_f, scan_name), 'jit_scan': jit_scan, 'jit_f': jit_f, 'scan': scan_impl} for jit_scan in [False, True] for jit_f in [False, True] for (scan_impl, scan_name) in SCAN_IMPLS))
def testScanLinearize(self, jit_scan, jit_f, scan):
    rng = np.random.RandomState(0)
    d = rng.randn(2)

    def f(c, a):
        assert (a.shape == (3,))
        assert (c.shape == (4,))
        b = jnp.cos(((jnp.sum(jnp.sin(a)) + jnp.sum(jnp.cos(c))) + jnp.sum(jnp.tan(d))))
        c = jnp.sin((c * b))
        assert (b.shape == ())
        return (c, b)
    if jit_f:
        f = api.jit(f)
    if jit_scan:
        scan = api.jit(scan, static_argnums=(0,))
    as_ = rng.randn(5, 3)
    c = rng.randn(4)
    ans = api.linearize((lambda c, as_: scan(f, c, as_)), c, as_)[1](c, as_)
    expected = api.linearize((lambda c, as_: scan_reference(f, c, as_)), c, as_)[1](c, as_)
    self.assertAllClose(ans, expected, check_dtypes=False, rtol={np.float64: 1e-14})