code: def test_custom_linear_solve_lu(self):
      def linear_solve(a, b):
          a_factors = jsp.linalg.lu_factor(a)
          at_factors = jsp.linalg.lu_factor(a.T)
  
          def solve(matvec, x):
              return jsp.linalg.lu_solve(a_factors, x)
  
          def transpose_solve(vecmat, x):
              return jsp.linalg.lu_solve(at_factors, x)
  
          return lax.custom_linear_solve(
              partial(high_precision_dot, a), b, solve, transpose_solve
          )
  
      rng = self.rng()
      a = rng.randn(3, 3)
      b = rng.randn(3)
  
      expected = jnp.linalg.solve(a, b)
      actual = linear_solve(a, b)
      self.assertAllClose(expected, actual)
  
      jtu.check_grads(linear_solve, (a, b), order=2, rtol=2e-3)
  
      # regression test for https://github.com/google/jax/issues/1536
      jtu.check_grads(
          jax.jit(linear_solve),
          (a, b),
          order=2,
          rtol={np.float32: 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 7 for calling the CustomLinearSolveTest obj's underlying member method code with example usage: obj.test_custom_linear_solve_lu() and returns: lax, jsp, jsp