def testCond(self):

    def fun(x):
        if (x < 3):
            return (x, x)
        else:
            y = lax.mul(2, x)
            return (y, lax.mul(2, y))

    @api.jit
    def cfun(x):

        def false_fun(x):
            y = lax.mul(2, x)
            return (y, lax.mul(2, y))
        return lax.cond(lax.lt(x, 3), (lambda x: (x, x)), false_fun, x)
    self.assertEqual(fun(0), cfun(0))
    self.assertEqual(fun(0), (0, 0))
    self.assertEqual(fun(1), cfun(1))
    self.assertEqual(fun(1), (1, 1))
    self.assertEqual(fun(2), cfun(2))
    self.assertEqual(fun(2), (2, 2))
    self.assertEqual(fun(3), cfun(3))
    self.assertEqual(fun(3), (6, 12))
    self.assertEqual(fun(4), cfun(4))
    self.assertEqual(fun(4), (8, 16))
def testCond(self):

    def fun(x):
        if (x < 3):
            return (x, x)
        else:
            y = lax.mul(2, x)
            return (y, lax.mul(2, y))

    @api.jit
    def cfun(x):

        def false_fun(x):
            y = lax.mul(2, x)
            return (y, lax.mul(2, y))
        return lax.cond(lax.lt(x, 3), (lambda x: (x, x)), false_fun, x)
    self.assertEqual(fun(0), cfun(0))
    self.assertEqual(fun(0), (0, 0))
    self.assertEqual(fun(1), cfun(1))
    self.assertEqual(fun(1), (1, 1))
    self.assertEqual(fun(2), cfun(2))
    self.assertEqual(fun(2), (2, 2))
    self.assertEqual(fun(3), cfun(3))
    self.assertEqual(fun(3), (6, 12))
    self.assertEqual(fun(4), cfun(4))
    self.assertEqual(fun(4), (8, 16))