code: def test_custom_root_vector_with_solve_closure(self):
      def vector_solve(f, y):
          return jnp.linalg.solve(jax.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 = self.rng()
      a = rng.randn(2, 2)
      b = rng.randn(2)
      jtu.check_grads(
          linear_solve,
          (a, b),
          order=2,
          atol={np.float32: 1e-2, np.float64: 1e-11},
      )
  
      actual = jax.jit(linear_solve)(a, b)
      expected = jnp.linalg.solve(a, b)
      self.assertAllClose(expected, actual)
  
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 2 for calling the CustomRootTest obj's underlying member method code with example usage: obj.test_custom_root_vector_with_solve_closure() and returns: jnp, lax