def testScanVmapFixpoint(self):

    def f(carry_init):

        def scan_body(c, x):
            return ((c[1], c[2], c[3], 0.0), None)
        return lax.scan(scan_body, (0.0, 1.0, 2.0, carry_init), jnp.zeros(2))
    carry_init = jnp.array([3.0, 4.0, 5.0])
    (carry_out, _) = api.vmap(f)(carry_init)
    self.assertAllClose(carry_out[3], jnp.array([0.0, 0.0, 0.0]), check_dtypes=False)
    self.assertAllClose(carry_out[2], jnp.array([0.0, 0.0, 0.0]), check_dtypes=False)
    self.assertAllClose(carry_out[1], carry_init, check_dtypes=False)
    self.assertAllClose(carry_out[0], jnp.array([2.0, 2.0, 2.0]), check_dtypes=False)
def testScanVmapFixpoint(self):

    def f(carry_init):

        def scan_body(c, x):
            return ((c[1], c[2], c[3], 0.0), None)
        return lax.scan(scan_body, (0.0, 1.0, 2.0, carry_init), jnp.zeros(2))
    carry_init = jnp.array([3.0, 4.0, 5.0])
    (carry_out, _) = api.vmap(f)(carry_init)
    self.assertAllClose(carry_out[3], jnp.array([0.0, 0.0, 0.0]), check_dtypes=False)
    self.assertAllClose(carry_out[2], jnp.array([0.0, 0.0, 0.0]), check_dtypes=False)
    self.assertAllClose(carry_out[1], carry_init, check_dtypes=False)
    self.assertAllClose(carry_out[0], jnp.array([2.0, 2.0, 2.0]), check_dtypes=False)