def test_scan_vjp_forwards_extensive_residuals(self):

    def cumprod(x):
        s = jnp.ones((2, 32), jnp.float32)
        return lax.scan((lambda s, x: ((x * s), s)), s, x)
    rng = np.random.RandomState(1234)
    x = jnp.asarray(rng.randn(32, 2, 32).astype('float32'))
    (_, vjp_fun) = api.vjp(cumprod, x)
    (*_, ext_res) = vjp_fun.args[0].args[0]
    self.assertIs(ext_res, x)
    x = rng.randn(32, 2, 32).astype('float32')
    (_, vjp_fun) = api.vjp(cumprod, x)
    (*_, ext_res) = vjp_fun.args[0].args[0]
    self.assertIsInstance(ext_res, xla.DeviceArray)
def test_scan_vjp_forwards_extensive_residuals(self):

    def cumprod(x):
        s = jnp.ones((2, 32), jnp.float32)
        return lax.scan((lambda s, x: ((x * s), s)), s, x)
    rng = np.random.RandomState(1234)
    x = jnp.asarray(rng.randn(32, 2, 32).astype('float32'))
    (_, vjp_fun) = api.vjp(cumprod, x)
    (*_, ext_res) = vjp_fun.args[0].args[0]
    self.assertIs(ext_res, x)
    x = rng.randn(32, 2, 32).astype('float32')
    (_, vjp_fun) = api.vjp(cumprod, x)
    (*_, ext_res) = vjp_fun.args[0].args[0]
    self.assertIsInstance(ext_res, xla.DeviceArray)