def testSwitch(self):

    def branch(x):
        y = lax.mul(2, x)
        return (y, lax.mul(2, y))
    branches = [(lambda x: (x, x)), branch, (lambda x: (x, (- x)))]

    def fun(x):
        if (x <= 0):
            return branches[0](x)
        elif (x == 1):
            return branches[1](x)
        else:
            return branches[2](x)

    def cfun(x):
        return lax.switch(x, branches, x)
    self.assertEqual(fun((- 1)), cfun((- 1)))
    self.assertEqual(fun(0), cfun(0))
    self.assertEqual(fun(1), cfun(1))
    self.assertEqual(fun(2), cfun(2))
    self.assertEqual(fun(3), cfun(3))
    cfun = api.jit(cfun)
    self.assertEqual(fun((- 1)), cfun((- 1)))
    self.assertEqual(fun(0), cfun(0))
    self.assertEqual(fun(1), cfun(1))
    self.assertEqual(fun(2), cfun(2))
    self.assertEqual(fun(3), cfun(3))
def testSwitch(self):

    def branch(x):
        y = lax.mul(2, x)
        return (y, lax.mul(2, y))
    branches = [(lambda x: (x, x)), branch, (lambda x: (x, (- x)))]

    def fun(x):
        if (x <= 0):
            return branches[0](x)
        elif (x == 1):
            return branches[1](x)
        else:
            return branches[2](x)

    def cfun(x):
        return lax.switch(x, branches, x)
    self.assertEqual(fun((- 1)), cfun((- 1)))
    self.assertEqual(fun(0), cfun(0))
    self.assertEqual(fun(1), cfun(1))
    self.assertEqual(fun(2), cfun(2))
    self.assertEqual(fun(3), cfun(3))
    cfun = api.jit(cfun)
    self.assertEqual(fun((- 1)), cfun((- 1)))
    self.assertEqual(fun(0), cfun(0))
    self.assertEqual(fun(1), cfun(1))
    self.assertEqual(fun(2), cfun(2))
    self.assertEqual(fun(3), cfun(3))