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

    def f(x):
        return cond((x < 2), (lambda x: (np.array([1.0, 2.0]) * x)), (lambda x: (np.array([3.0, 4.0]) * jnp.sin(x))), x)

    def f_ref(x):
        if (x < 2):
            return (np.array([1.0, 2.0]) * x)
        else:
            return (np.array([3.0, 4.0]) * np.sin(x))
    y = f(1.0)
    expected = f_ref(1.0)
    self.assertAllClose(y, expected, check_dtypes=False)
    y = f(4.0)
    expected = f_ref(4.0)
    self.assertAllClose(y, expected, check_dtypes=False)
@parameterized.named_parameters(({'testcase_name': f'_{name}', 'cond': cond} for (cond, name) in COND_IMPLS))
def testCondWithConsts(self, cond):

    def f(x):
        return cond((x < 2), (lambda x: (np.array([1.0, 2.0]) * x)), (lambda x: (np.array([3.0, 4.0]) * jnp.sin(x))), x)

    def f_ref(x):
        if (x < 2):
            return (np.array([1.0, 2.0]) * x)
        else:
            return (np.array([3.0, 4.0]) * np.sin(x))
    y = f(1.0)
    expected = f_ref(1.0)
    self.assertAllClose(y, expected, check_dtypes=False)
    y = f(4.0)
    expected = f_ref(4.0)
    self.assertAllClose(y, expected, check_dtypes=False)