def testCondJit(self):

    def f(x):
        return lax.cond((x < 2), (lambda x: (3.0 * x)), (lambda x: jnp.sin(x)), x)
    y = api.jit(f)(1.0)
    expected = f(1.0)
    self.assertAllClose(y, expected, check_dtypes=False)
    y = api.jit(f)(4.0)
    expected = f(4.0)
    self.assertAllClose(y, expected, check_dtypes=False)
def testCondJit(self):

    def f(x):
        return lax.cond((x < 2), (lambda x: (3.0 * x)), (lambda x: jnp.sin(x)), x)
    y = api.jit(f)(1.0)
    expected = f(1.0)
    self.assertAllClose(y, expected, check_dtypes=False)
    y = api.jit(f)(4.0)
    expected = f(4.0)
    self.assertAllClose(y, expected, check_dtypes=False)