code: 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 = self.rng()
      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 = jax.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: 1e-2},
      )
  
imports: import re
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import jax
  from jax import lax
  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 CustomRootTest obj's underlying member method code with example usage: obj.test_custom_root_with_custom_linear_solve() and returns: lax, lax