def testCondLinearize(self):

    def f(x):
        return lax.cond((x < 2), (lambda x: (3.0 * x)), (lambda x: jnp.sin(x)), x)
    (y, f_lin) = api.linearize(f, 1.0)
    self.assertAllClose(y, 3.0, check_dtypes=False)
    self.assertAllClose(f_lin(2.0), 6.0, check_dtypes=False)
    (y, f_lin) = api.linearize(f, 4.0)
    self.assertAllClose(y, jnp.sin(4.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), (jnp.cos(4.0) * 2.0), check_dtypes=False)
def testCondLinearize(self):

    def f(x):
        return lax.cond((x < 2), (lambda x: (3.0 * x)), (lambda x: jnp.sin(x)), x)
    (y, f_lin) = api.linearize(f, 1.0)
    self.assertAllClose(y, 3.0, check_dtypes=False)
    self.assertAllClose(f_lin(2.0), 6.0, check_dtypes=False)
    (y, f_lin) = api.linearize(f, 4.0)
    self.assertAllClose(y, jnp.sin(4.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), (jnp.cos(4.0) * 2.0), check_dtypes=False)