def testScanVmapTuples(self):

    def f(c, a):
        (a1, a2) = a
        (c1, c2) = c
        b = (jnp.sum(jnp.cos(a1)) * jnp.sum(jnp.tan((c2 * a2))))
        c = ((c1 * jnp.sin(jnp.sum((a1 * a2)))), (c2 * jnp.cos(jnp.sum(a1))))
        return (c, b)
    in_axes = (0, (1, 2))
    r = np.random.RandomState(0)
    as_ = (r.randn(3, 7), r.randn(3, 4, 7))
    c = (r.randn(7, 2), r.randn(7))
    (expected_c_out, expected_bs) = ([], [])
    for i in range(7):
        (c_out, bs) = lax.scan(f, (c[0][i], c[1][i]), (as_[0][:, i], as_[1][:, :, i]))
        expected_c_out.append(c_out)
        expected_bs.append(bs)
    (expected_c_out_0, expected_c_out_1) = unzip2(expected_c_out)
    expected_c_out = (jnp.stack(expected_c_out_0), jnp.stack(expected_c_out_1))
    expected_bs = jnp.stack(expected_bs)
    expected = (expected_c_out, expected_bs)
    ans = api.vmap((lambda c, as_: lax.scan(f, c, as_)), in_axes)(c, as_)
    self.assertAllClose(ans, expected, check_dtypes=False)
def testScanVmapTuples(self):

    def f(c, a):
        (a1, a2) = a
        (c1, c2) = c
        b = (jnp.sum(jnp.cos(a1)) * jnp.sum(jnp.tan((c2 * a2))))
        c = ((c1 * jnp.sin(jnp.sum((a1 * a2)))), (c2 * jnp.cos(jnp.sum(a1))))
        return (c, b)
    in_axes = (0, (1, 2))
    r = np.random.RandomState(0)
    as_ = (r.randn(3, 7), r.randn(3, 4, 7))
    c = (r.randn(7, 2), r.randn(7))
    (expected_c_out, expected_bs) = ([], [])
    for i in range(7):
        (c_out, bs) = lax.scan(f, (c[0][i], c[1][i]), (as_[0][:, i], as_[1][:, :, i]))
        expected_c_out.append(c_out)
        expected_bs.append(bs)
    (expected_c_out_0, expected_c_out_1) = unzip2(expected_c_out)
    expected_c_out = (jnp.stack(expected_c_out_0), jnp.stack(expected_c_out_1))
    expected_bs = jnp.stack(expected_bs)
    expected = (expected_c_out, expected_bs)
    ans = api.vmap((lambda c, as_: lax.scan(f, c, as_)), in_axes)(c, as_)
    self.assertAllClose(ans, expected, check_dtypes=False)