def testCondJVP(self):

    def fun_ref(x):
        if (x < 3):
            return (x, x)
        else:
            y = (2 * x)
            return (y, (2 * y))

    def fun(x):

        def false_fun(x):
            y = (2 * x)
            return (y, (2 * y))
        return lax.cond((x < 3), (lambda x: (x, x)), false_fun, x)
    x = 3.14
    ans = api.jvp(fun, (x,), (x,))
    expected = api.jvp(fun_ref, (x,), (x,))
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(fun, (x,), order=2, modes=['fwd'])
    x = 2.72
    ans = api.jvp(fun, (x,), (x,))
    expected = api.jvp(fun_ref, (x,), (x,))
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(fun, (x,), order=2, modes=['fwd'])
def testCondJVP(self):

    def fun_ref(x):
        if (x < 3):
            return (x, x)
        else:
            y = (2 * x)
            return (y, (2 * y))

    def fun(x):

        def false_fun(x):
            y = (2 * x)
            return (y, (2 * y))
        return lax.cond((x < 3), (lambda x: (x, x)), false_fun, x)
    x = 3.14
    ans = api.jvp(fun, (x,), (x,))
    expected = api.jvp(fun_ref, (x,), (x,))
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(fun, (x,), order=2, modes=['fwd'])
    x = 2.72
    ans = api.jvp(fun, (x,), (x,))
    expected = api.jvp(fun_ref, (x,), (x,))
    self.assertAllClose(ans, expected, check_dtypes=False)
    jtu.check_grads(fun, (x,), order=2, modes=['fwd'])