def testWhileLoopBatched(self):

    def fun(x):
        return lax.while_loop((lambda x: (x < 3)), (lambda x: (x + 2)), x)
    ans = api.vmap(fun)(np.array([0, 1, 2, 3]))
    expected = np.array([4, 3, 4, 3])
    self.assertAllClose(ans, expected, check_dtypes=False)
    fun = api.jit(fun)
    ans = api.vmap(fun)(np.array([0, 1, 2, 3]))
    expected = np.array([4, 3, 4, 3])
    self.assertAllClose(ans, expected, check_dtypes=False)
def testWhileLoopBatched(self):

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