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

    def f_ref(x):
        z = ((jnp.array([1.0, 2.0]) * x) if (x[0] < 2) else jnp.cos(jnp.sin(x)))
        return z.sum()

    def f(x):
        return cond((x[0] < 2), (lambda x: (jnp.array([1.0, 2.0]) * x)), (lambda x: jnp.cos(jnp.sin(x))), x).sum()
    x = (2.14 * jnp.ones(2))
    (y, f_lin) = api.linearize(f, x)
    (y_ref, f_lin_ref) = api.linearize(f_ref, x)
    self.assertAllClose(y, y_ref, check_dtypes=False)
    self.assertAllClose(f_lin(x), f_lin_ref(x), check_dtypes=False)
    x = ((- 2.14) * jnp.ones(2))
    (y, f_lin) = api.linearize(f, x)
    (y_ref, f_lin_ref) = api.linearize(f_ref, x)
    self.assertAllClose(y, y_ref, check_dtypes=False)
    self.assertAllClose(f_lin(x), f_lin_ref(x), check_dtypes=False)
    f = api.jit(f)
    x = (2.14 * jnp.ones(2))
    (y, f_lin) = api.linearize(f, x)
    (y_ref, f_lin_ref) = api.linearize(f_ref, x)
    self.assertAllClose(y, y_ref, check_dtypes=False)
    self.assertAllClose(f_lin(x), f_lin_ref(x), check_dtypes=False)
@parameterized.named_parameters(({'testcase_name': f'_{name}', 'cond': cond} for (cond, name) in COND_IMPLS))
def testCondLinearize2(self, cond):

    def f_ref(x):
        z = ((jnp.array([1.0, 2.0]) * x) if (x[0] < 2) else jnp.cos(jnp.sin(x)))
        return z.sum()

    def f(x):
        return cond((x[0] < 2), (lambda x: (jnp.array([1.0, 2.0]) * x)), (lambda x: jnp.cos(jnp.sin(x))), x).sum()
    x = (2.14 * jnp.ones(2))
    (y, f_lin) = api.linearize(f, x)
    (y_ref, f_lin_ref) = api.linearize(f_ref, x)
    self.assertAllClose(y, y_ref, check_dtypes=False)
    self.assertAllClose(f_lin(x), f_lin_ref(x), check_dtypes=False)
    x = ((- 2.14) * jnp.ones(2))
    (y, f_lin) = api.linearize(f, x)
    (y_ref, f_lin_ref) = api.linearize(f_ref, x)
    self.assertAllClose(y, y_ref, check_dtypes=False)
    self.assertAllClose(f_lin(x), f_lin_ref(x), check_dtypes=False)
    f = api.jit(f)
    x = (2.14 * jnp.ones(2))
    (y, f_lin) = api.linearize(f, x)
    (y_ref, f_lin_ref) = api.linearize(f_ref, x)
    self.assertAllClose(y, y_ref, check_dtypes=False)
    self.assertAllClose(f_lin(x), f_lin_ref(x), check_dtypes=False)