code: def test_custom_linear_solve(self, symmetric):
      def explicit_jacobian_solve(matvec, b):
          return lax.stop_gradient(
              jnp.linalg.solve(jax.jacobian(matvec)(b), b)
          )
  
      def matrix_free_solve(matvec, b):
          return lax.custom_linear_solve(
              matvec,
              b,
              explicit_jacobian_solve,
              explicit_jacobian_solve,
              symmetric=symmetric,
          )
  
      def linear_solve(a, b):
          return matrix_free_solve(partial(high_precision_dot, a), b)
  
      rng = self.rng()
      a = rng.randn(3, 3)
      if symmetric:
          a = a + a.T
      b = rng.randn(3)
      jtu.check_grads(linear_solve, (a, b), order=2, rtol=3e-3)
  
      expected = jnp.linalg.solve(a, b)
      actual = jax.jit(linear_solve)(a, b)
      self.assertAllClose(expected, actual)
  
      c = rng.randn(3, 2)
      expected = jnp.linalg.solve(a, c)
      actual = jax.vmap(linear_solve, (None, 1), 1)(a, c)
      self.assertAllClose(expected, actual)
  
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 1 for calling the CustomLinearSolveTest obj's underlying member method code with example usage: obj.test_custom_linear_solve(symmetric) and returns: lax, lax, matrix_free_solve