def test_custom_root_with_custom_linear_solve(self):

    def linear_solve(a, b):
        f = (lambda x: (high_precision_dot(a, x) - b))
        factors = jsp.linalg.cho_factor(a)
        cho_solve = (lambda f, b: jsp.linalg.cho_solve(factors, b))

        def pos_def_solve(g, b):
            return lax.custom_linear_solve(g, b, cho_solve, symmetric=True)
        return lax.custom_root(f, b, cho_solve, pos_def_solve)
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    actual = linear_solve(high_precision_dot(a, a.T), b)
    expected = jnp.linalg.solve(high_precision_dot(a, a.T), b)
    self.assertAllClose(expected, actual)
    actual = api.jit(linear_solve)(high_precision_dot(a, a.T), b)
    expected = jnp.linalg.solve(high_precision_dot(a, a.T), b)
    self.assertAllClose(expected, actual)
    jtu.check_grads((lambda x, y: linear_solve(high_precision_dot(x, x.T), y)), (a, b), order=2, rtol={jnp.float32: 0.01})
def test_custom_root_with_custom_linear_solve(self):

    def linear_solve(a, b):
        f = (lambda x: (high_precision_dot(a, x) - b))
        factors = jsp.linalg.cho_factor(a)
        cho_solve = (lambda f, b: jsp.linalg.cho_solve(factors, b))

        def pos_def_solve(g, b):
            return lax.custom_linear_solve(g, b, cho_solve, symmetric=True)
        return lax.custom_root(f, b, cho_solve, pos_def_solve)
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    actual = linear_solve(high_precision_dot(a, a.T), b)
    expected = jnp.linalg.solve(high_precision_dot(a, a.T), b)
    self.assertAllClose(expected, actual)
    actual = api.jit(linear_solve)(high_precision_dot(a, a.T), b)
    expected = jnp.linalg.solve(high_precision_dot(a, a.T), b)
    self.assertAllClose(expected, actual)
    jtu.check_grads((lambda x, y: linear_solve(high_precision_dot(x, x.T), y)), (a, b), order=2, rtol={jnp.float32: 0.01})