@parameterized.named_parameters(({'testcase_name': '_jit_scan={}_jit_f={}_in_axes={}_impl={}'.format(jit_scan, jit_f, in_axes, scan_name), 'jit_scan': jit_scan, 'jit_f': jit_f, 'in_axes': in_axes, 'scan': scan_impl} for jit_scan in [False, True] for jit_f in [False, True] for (scan_impl, scan_name) in SCAN_IMPLS for in_axes in itertools.product([None, 0, 1], [None, 0, 1, 2]) if (in_axes != (None, None))))
def testScanVmap(self, jit_scan, jit_f, in_axes, 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_shape = [5, 3]
    c_shape = [4]
    (c_bdim, as_bdim) = in_axes
    if (c_bdim is not None):
        c_shape.insert(c_bdim, 7)
    if (as_bdim is not None):
        as_shape.insert(as_bdim, 7)
    as_ = rng.randn(*as_shape)
    c = rng.randn(*c_shape)
    ans = api.vmap((lambda c, as_: scan(f, c, as_)), in_axes)(c, as_)
    expected = api.vmap((lambda c, as_: scan_reference(f, c, as_)), in_axes)(c, as_)
    self.assertAllClose(ans, expected, check_dtypes=False, rtol=1e-05, atol=1e-05)
@parameterized.named_parameters(({'testcase_name': '_jit_scan={}_jit_f={}_in_axes={}_impl={}'.format(jit_scan, jit_f, in_axes, scan_name), 'jit_scan': jit_scan, 'jit_f': jit_f, 'in_axes': in_axes, 'scan': scan_impl} for jit_scan in [False, True] for jit_f in [False, True] for (scan_impl, scan_name) in SCAN_IMPLS for in_axes in itertools.product([None, 0, 1], [None, 0, 1, 2]) if (in_axes != (None, None))))
def testScanVmap(self, jit_scan, jit_f, in_axes, 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_shape = [5, 3]
    c_shape = [4]
    (c_bdim, as_bdim) = in_axes
    if (c_bdim is not None):
        c_shape.insert(c_bdim, 7)
    if (as_bdim is not None):
        as_shape.insert(as_bdim, 7)
    as_ = rng.randn(*as_shape)
    c = rng.randn(*c_shape)
    ans = api.vmap((lambda c, as_: scan(f, c, as_)), in_axes)(c, as_)
    expected = api.vmap((lambda c, as_: scan_reference(f, c, as_)), in_axes)(c, as_)
    self.assertAllClose(ans, expected, check_dtypes=False, rtol=1e-05, atol=1e-05)