def testForiLoopBatched(self):

    def body_fun(i, loop_carry):
        (x, y) = loop_carry
        x = (x + 1)
        y = (y + 2)
        return (x, y)

    def fun(x):
        return lax.fori_loop(0, 10, body_fun, (x, 0))
    ans = api.vmap(fun)(np.array([0, 1]))
    expected = (np.array([10, 11]), np.array([20, 20]))
    self.assertAllClose(ans, expected, check_dtypes=False)
def testForiLoopBatched(self):

    def body_fun(i, loop_carry):
        (x, y) = loop_carry
        x = (x + 1)
        y = (y + 2)
        return (x, y)

    def fun(x):
        return lax.fori_loop(0, 10, body_fun, (x, 0))
    ans = api.vmap(fun)(np.array([0, 1]))
    expected = (np.array([10, 11]), np.array([20, 20]))
    self.assertAllClose(ans, expected, check_dtypes=False)