def testSwitchGrad(self):
    branches = [(lambda x: (3.0 * x)), (lambda x: jnp.sin(x)), (lambda x: (- x))]

    def f_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 f(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.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 testSwitchGrad(self):
    branches = [(lambda x: (3.0 * x)), (lambda x: jnp.sin(x)), (lambda x: (- x))]

    def f_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 f(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.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'])