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

    def richardson_iteration(matvec, b, omega=0.1, tolerance=1e-06):

        def cond(x):
            return (jnp.linalg.norm((matvec(x) - b)) > tolerance)

        def body(x):
            return (x + (omega * (b - matvec(x))))
        return lax.while_loop(cond, body, b)

    def matrix_free_solve(matvec, b):
        return lax.custom_linear_solve(matvec, b, richardson_iteration, richardson_iteration)

    def build_and_solve(a, b):
        matvec = partial(high_precision_dot, jnp.exp(a))
        return matrix_free_solve(matvec, jnp.cos(b))
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    expected = jnp.linalg.solve(jnp.exp(a), jnp.cos(b))
    actual = build_and_solve(a, b)
    self.assertAllClose(expected, actual, atol=1e-05)
    jtu.check_grads(build_and_solve, (a, b), atol=1e-05, order=2, rtol={jnp.float32: 0.06, jnp.float64: 0.002})
    jtu.check_grads(api.vmap(build_and_solve), (a[None, :, :], b[None, :]), atol=1e-05, order=2, rtol={jnp.float32: 0.06, jnp.float64: 0.002})
@jtu.skip_on_flag('jax_skip_slow_tests', True)
def test_custom_linear_solve_iterative(self):

    def richardson_iteration(matvec, b, omega=0.1, tolerance=1e-06):

        def cond(x):
            return (jnp.linalg.norm((matvec(x) - b)) > tolerance)

        def body(x):
            return (x + (omega * (b - matvec(x))))
        return lax.while_loop(cond, body, b)

    def matrix_free_solve(matvec, b):
        return lax.custom_linear_solve(matvec, b, richardson_iteration, richardson_iteration)

    def build_and_solve(a, b):
        matvec = partial(high_precision_dot, jnp.exp(a))
        return matrix_free_solve(matvec, jnp.cos(b))
    rng = np.random.RandomState(0)
    a = rng.randn(2, 2)
    b = rng.randn(2)
    expected = jnp.linalg.solve(jnp.exp(a), jnp.cos(b))
    actual = build_and_solve(a, b)
    self.assertAllClose(expected, actual, atol=1e-05)
    jtu.check_grads(build_and_solve, (a, b), atol=1e-05, order=2, rtol={jnp.float32: 0.06, jnp.float64: 0.002})
    jtu.check_grads(api.vmap(build_and_solve), (a[None, :, :], b[None, :]), atol=1e-05, order=2, rtol={jnp.float32: 0.06, jnp.float64: 0.002})