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

    def f_ref(x):
        z = ((jnp.array([1.0, 2.0]) * x) if (x[0] < 2) else 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.sin(x)), x)
    f = (lambda x: api.jit(_f)(x).sum())
    x = (2.14 * jnp.ones(2))
    ans = api.grad(f)(x)
    expected = api.grad(f_ref)(x)
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(f, (x,), order=2, modes=['fwd', 'rev'])
    x = (1.72 * jnp.ones(2))
    ans = api.grad(f)(x)
    expected = api.grad(f_ref)(x)
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(f, (x,), order=2, modes=['fwd', 'rev'], rtol={jnp.float32: 0.01, jnp.float64: 0.002})
@parameterized.named_parameters(({'testcase_name': f'_{name}', 'cond': cond} for (cond, name) in COND_IMPLS))
def testCondGrad2(self, cond):

    def f_ref(x):
        z = ((jnp.array([1.0, 2.0]) * x) if (x[0] < 2) else 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.sin(x)), x)
    f = (lambda x: api.jit(_f)(x).sum())
    x = (2.14 * jnp.ones(2))
    ans = api.grad(f)(x)
    expected = api.grad(f_ref)(x)
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(f, (x,), order=2, modes=['fwd', 'rev'])
    x = (1.72 * jnp.ones(2))
    ans = api.grad(f)(x)
    expected = api.grad(f_ref)(x)
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(f, (x,), order=2, modes=['fwd', 'rev'], rtol={jnp.float32: 0.01, jnp.float64: 0.002})