def testSwitchJVP(self):

    def branch(x):
        y = (2 * x)
        return (y, (2 * y))
    branches = [(lambda x: (x, x)), branch, (lambda x: (x, (- x)))]

    def fun_ref(x):
        idx = (x // 1)
        if (idx <= 0):
            return branches[0](x)
        elif (idx == 1):
            return branches[1](x)
        else:
            return branches[2](x)

    def fun(x):
        idx = lax.convert_element_type((x // 1), np.int32)
        return lax.switch(idx, branches, x)
    for x in [(- 0.7), 0.7, 1.7, 2.7, 3.7]:
        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 testSwitchJVP(self):

    def branch(x):
        y = (2 * x)
        return (y, (2 * y))
    branches = [(lambda x: (x, x)), branch, (lambda x: (x, (- x)))]

    def fun_ref(x):
        idx = (x // 1)
        if (idx <= 0):
            return branches[0](x)
        elif (idx == 1):
            return branches[1](x)
        else:
            return branches[2](x)

    def fun(x):
        idx = lax.convert_element_type((x // 1), np.int32)
        return lax.switch(idx, branches, x)
    for x in [(- 0.7), 0.7, 1.7, 2.7, 3.7]:
        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'])