def testOneBranchSwitch(self):
    branch = (lambda x: (- x))
    f = (lambda i, x: lax.switch(i, [branch], x))
    x = 7.0
    self.assertEqual(f((- 1), x), branch(x))
    self.assertEqual(f(0, x), branch(x))
    self.assertEqual(f(1, x), branch(x))
    cf = api.jit(f)
    self.assertEqual(cf((- 1), x), branch(x))
    self.assertEqual(cf(0, x), branch(x))
    self.assertEqual(cf(1, x), branch(x))
    cf = api.jit(f, static_argnums=0)
    self.assertEqual(cf((- 1), x), branch(x))
    self.assertEqual(cf(0, x), branch(x))
    self.assertEqual(cf(1, x), branch(x))
def testOneBranchSwitch(self):
    branch = (lambda x: (- x))
    f = (lambda i, x: lax.switch(i, [branch], x))
    x = 7.0
    self.assertEqual(f((- 1), x), branch(x))
    self.assertEqual(f(0, x), branch(x))
    self.assertEqual(f(1, x), branch(x))
    cf = api.jit(f)
    self.assertEqual(cf((- 1), x), branch(x))
    self.assertEqual(cf(0, x), branch(x))
    self.assertEqual(cf(1, x), branch(x))
    cf = api.jit(f, static_argnums=0)
    self.assertEqual(cf((- 1), x), branch(x))
    self.assertEqual(cf(0, x), branch(x))
    self.assertEqual(cf(1, x), branch(x))