code: def test_custom_linear_solve_iterative(self):
      def richardson_iteration(matvec, b, omega=0.1, tolerance=1e-6):
          # Equivalent to vanilla gradient descent:
          # https://en.wikipedia.org/wiki/Modified_Richardson_iteration
          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):
          # intentionally non-linear in a and b
          matvec = partial(high_precision_dot, jnp.exp(a))
          return matrix_free_solve(matvec, jnp.cos(b))
  
      # rng = self.rng()
      # This test is very sensitive to the inputs, so we use a known working seed.
      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-5)
      jtu.check_grads(
          build_and_solve,
          (a, b),
          atol=1e-5,
          order=2,
          rtol={jnp.float32: 6e-2, jnp.float64: 2e-3},
      )
  
      # vmap across an empty dimension
      jtu.check_grads(
          jax.vmap(build_and_solve),
          (a[None, :, :], b[None, :]),
          atol=1e-5,
          order=2,
          rtol={jnp.float32: 6e-2, jnp.float64: 2e-3},
      )
  
imports: from functools import partial
  import re
  import unittest
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import jax
  from jax import lax
  from jax.ad_checkpoint import checkpoint
  from jax._src import test_util
  from jax import tree_util
  import jax.numpy
  import jax.scipy
  from jax import config
label_desc: Use image node_id 4 for calling the CustomLinearSolveTest obj's underlying member method code with example usage: obj.test_custom_linear_solve_iterative() and returns: lax, lax, matrix_free_solve, unknown, unknown