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

    def f(x):
        idx = lax.convert_element_type((x // 1), np.int32)
        return lax.switch(idx, branches, x)
    for x in [(- 1.0), 0.0, 1.0, 2.0, 3.0]:
        y = api.jit(f)(x)
        expected = f(x)
        self.assertAllClose(y, expected, check_dtypes=False)
def testSwitchJit(self):
    branches = [(lambda x: (3.0 * x)), (lambda x: jnp.sin(x)), (lambda x: (- x))]

    def f(x):
        idx = lax.convert_element_type((x // 1), np.int32)
        return lax.switch(idx, branches, x)
    for x in [(- 1.0), 0.0, 1.0, 2.0, 3.0]:
        y = api.jit(f)(x)
        expected = f(x)
        self.assertAllClose(y, expected, check_dtypes=False)