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

    def fun_ref(x, y):
        if (x < 3):
            return (2.0 * jnp.sin(y))
        else:
            return (2.0 * jnp.cos(x))

    def fun(x, y):
        return cond((x < 3), (), (lambda _: (2.0 * jnp.sin(y))), x, (lambda x: (2.0 * x)))
    y = 5.8
    x = 3.14
    ans = api.grad(fun, 1)(x, y)
    expected = api.grad(fun_ref, 1)(x, y)
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(fun, (x, y), order=2, modes=['fwd', 'rev'])
    x = 2.72
    ans = api.grad(fun, 1)(x, y)
    expected = api.grad(fun_ref, 1)(x, y)
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(fun, (x, y), order=2, modes=['fwd', 'rev'])
@parameterized.named_parameters(({'testcase_name': f'_{name}', 'cond': cond} for (cond, name) in COND_IMPLS))
def testCondGrad4(self, cond):

    def fun_ref(x, y):
        if (x < 3):
            return (2.0 * jnp.sin(y))
        else:
            return (2.0 * jnp.cos(x))

    def fun(x, y):
        return cond((x < 3), (), (lambda _: (2.0 * jnp.sin(y))), x, (lambda x: (2.0 * x)))
    y = 5.8
    x = 3.14
    ans = api.grad(fun, 1)(x, y)
    expected = api.grad(fun_ref, 1)(x, y)
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(fun, (x, y), order=2, modes=['fwd', 'rev'])
    x = 2.72
    ans = api.grad(fun, 1)(x, y)
    expected = api.grad(fun_ref, 1)(x, y)
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(fun, (x, y), order=2, modes=['fwd', 'rev'])