@jtu.skip_on_flag('jax_skip_slow_tests', True)
def test_custom_root_vector_with_solve_closure(self):

    def vector_solve(f, y):
        return jnp.linalg.solve(api.jacobian(f)(y), y)

    def linear_solve(a, b):
        f = (lambda y: (high_precision_dot(a, y) - b))
        x0 = jnp.zeros_like(b)
        solution = jnp.linalg.solve(a, b)
        oracle = (lambda func, x0: solution)
        return lax.custom_root(f, x0, oracle, vector_solve)
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    jtu.check_grads(linear_solve, (a, b), order=2, atol={np.float32: 0.01, np.float64: 1e-11})
    actual = api.jit(linear_solve)(a, b)
    expected = jnp.linalg.solve(a, b)
    self.assertAllClose(expected, actual)
@jtu.skip_on_flag('jax_skip_slow_tests', True)
def test_custom_root_vector_with_solve_closure(self):

    def vector_solve(f, y):
        return jnp.linalg.solve(api.jacobian(f)(y), y)

    def linear_solve(a, b):
        f = (lambda y: (high_precision_dot(a, y) - b))
        x0 = jnp.zeros_like(b)
        solution = jnp.linalg.solve(a, b)
        oracle = (lambda func, x0: solution)
        return lax.custom_root(f, x0, oracle, vector_solve)
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    jtu.check_grads(linear_solve, (a, b), order=2, atol={np.float32: 0.01, np.float64: 1e-11})
    actual = api.jit(linear_solve)(a, b)
    expected = jnp.linalg.solve(a, b)
    self.assertAllClose(expected, actual)