def testCondGrad(self):

    def f_ref(x):
        return ((3.0 * x) if (x < 2) else jnp.sin(x))

    def f(x):
        return lax.cond((x < 2), (lambda x: (3.0 * x)), (lambda x: jnp.sin(x)), x)
    x = 2.14
    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
    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'])
def testCondGrad(self):

    def f_ref(x):
        return ((3.0 * x) if (x < 2) else jnp.sin(x))

    def f(x):
        return lax.cond((x < 2), (lambda x: (3.0 * x)), (lambda x: jnp.sin(x)), x)
    x = 2.14
    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
    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'])