def testSwitchLinearize(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)
    (y, f_lin) = api.linearize(f, (- 1.0))
    self.assertAllClose(y, (- 3.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), 6.0, check_dtypes=False)
    (y, f_lin) = api.linearize(f, 0.0)
    self.assertAllClose(y, 0.0, check_dtypes=False)
    self.assertAllClose(f_lin(2.0), 6.0, check_dtypes=False)
    (y, f_lin) = api.linearize(f, 1.0)
    self.assertAllClose(y, jnp.sin(1.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), (jnp.cos(1.0) * 2.0), check_dtypes=False)
    (y, f_lin) = api.linearize(f, 2.0)
    self.assertAllClose(y, (- 2.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), (- 2.0), check_dtypes=False)
    (y, f_lin) = api.linearize(f, 3.0)
    self.assertAllClose(y, (- 3.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), (- 2.0), check_dtypes=False)
def testSwitchLinearize(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)
    (y, f_lin) = api.linearize(f, (- 1.0))
    self.assertAllClose(y, (- 3.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), 6.0, check_dtypes=False)
    (y, f_lin) = api.linearize(f, 0.0)
    self.assertAllClose(y, 0.0, check_dtypes=False)
    self.assertAllClose(f_lin(2.0), 6.0, check_dtypes=False)
    (y, f_lin) = api.linearize(f, 1.0)
    self.assertAllClose(y, jnp.sin(1.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), (jnp.cos(1.0) * 2.0), check_dtypes=False)
    (y, f_lin) = api.linearize(f, 2.0)
    self.assertAllClose(y, (- 2.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), (- 2.0), check_dtypes=False)
    (y, f_lin) = api.linearize(f, 3.0)
    self.assertAllClose(y, (- 3.0), check_dtypes=False)
    self.assertAllClose(f_lin(2.0), (- 2.0), check_dtypes=False)