def testWhileLoopTupleBatched(self):

    def cond_fun(loop_carry):
        (x, y) = loop_carry
        return ((x + y) < 5)

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

    def fun(x, y):
        return lax.while_loop(cond_fun, body_fun, (x, y))
    ans = api.vmap(fun)(np.array([0, 0]), np.array([1, 2]))
    expected = (np.array([4, 3]), np.array([1, 2]))
    self.assertAllClose(ans, expected, check_dtypes=False)
def testWhileLoopTupleBatched(self):

    def cond_fun(loop_carry):
        (x, y) = loop_carry
        return ((x + y) < 5)

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

    def fun(x, y):
        return lax.while_loop(cond_fun, body_fun, (x, y))
    ans = api.vmap(fun)(np.array([0, 0]), np.array([1, 2]))
    expected = (np.array([4, 3]), np.array([1, 2]))
    self.assertAllClose(ans, expected, check_dtypes=False)