code: def test_custom_linear_solve_without_transpose_solve(self):
      def explicit_jacobian_solve(matvec, b):
          return lax.stop_gradient(
              jnp.linalg.solve(jax.jacobian(matvec)(b), b)
          )
  
      def loss(a, b):
          matvec = partial(high_precision_dot, a)
          x = lax.custom_linear_solve(
              matvec, b, explicit_jacobian_solve
          )
          return jnp.sum(x)
  
      rng = self.rng()
      a = rng.randn(2, 2)
      b = rng.randn(2)
  
      jtu.check_grads(
          loss,
          (a, b),
          order=2,
          modes=["fwd"],
          atol={np.float32: 2e-3, np.float64: 1e-11},
      )
      jtu.check_grads(
          jax.vmap(loss),
          (a[None, :, :], b[None, :]),
          order=2,
          modes=["fwd"],
          atol={np.float32: 2e-3, np.float64: 1e-11},
      )
  
      with self.assertRaisesRegex(
          TypeError, "transpose_solve required"
      ):
          jax.grad(loss)(a, b)
  
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 8 for calling the CustomLinearSolveTest obj's underlying member method code with example usage: obj.test_custom_linear_solve_without_transpose_solve() and returns: lax, jnp