@parameterized.named_parameters(({'testcase_name': f'_{name}', 'cond': cond} for (cond, name) in COND_IMPLS))
def testNestedCond(self, cond):

    def fun(x):
        if (x < 2):
            return lax.mul(2, x)
        elif (x < 5):
            return lax.mul(3, x)
        else:
            return lax.mul(4, x)

    @api.jit
    def cfun(x):
        return cond(lax.lt(x, 2), (lambda x: lax.mul(2, x)), (lambda x: cond(lax.lt(x, 5), x, (lambda x: lax.mul(3, x)), 4, (lambda y: lax.mul(y, x)))), x)
    self.assertEqual(cfun(1), 2)
    self.assertEqual(cfun(3), 9)
    self.assertEqual(cfun(6), 24)
    self.assertEqual(cfun(1), fun(1))
    self.assertEqual(cfun(3), fun(3))
    self.assertEqual(cfun(6), fun(6))
@parameterized.named_parameters(({'testcase_name': f'_{name}', 'cond': cond} for (cond, name) in COND_IMPLS))
def testNestedCond(self, cond):

    def fun(x):
        if (x < 2):
            return lax.mul(2, x)
        elif (x < 5):
            return lax.mul(3, x)
        else:
            return lax.mul(4, x)

    @api.jit
    def cfun(x):
        return cond(lax.lt(x, 2), (lambda x: lax.mul(2, x)), (lambda x: cond(lax.lt(x, 5), x, (lambda x: lax.mul(3, x)), 4, (lambda y: lax.mul(y, x)))), x)
    self.assertEqual(cfun(1), 2)
    self.assertEqual(cfun(3), 9)
    self.assertEqual(cfun(6), 24)
    self.assertEqual(cfun(1), fun(1))
    self.assertEqual(cfun(3), fun(3))
    self.assertEqual(cfun(6), fun(6))