code: def test_custom_linear_solve_cholesky(self):
      def positive_definite_solve(a, b):
          factors = jsp.linalg.cho_factor(a)
  
          def solve(matvec, x):
              return jsp.linalg.cho_solve(factors, x)
  
          matvec = partial(high_precision_dot, a)
          return lax.custom_linear_solve(
              matvec, b, solve, symmetric=True
          )
  
      rng = self.rng()
      a = rng.randn(2, 2)
      b = rng.randn(2)
  
      tol = {
          np.float32: 1e-3
          if jtu.test_device_matches(["tpu"])
          else 2e-4,
          np.float64: 1e-12,
      }
      expected = jnp.linalg.solve(np.asarray(posify(a)), b)
      actual = positive_definite_solve(posify(a), b)
      self.assertAllClose(expected, actual, rtol=tol, atol=tol)
  
      actual = jax.jit(positive_definite_solve)(posify(a), b)
      self.assertAllClose(expected, actual, rtol=tol, atol=tol)
  
      # numerical gradients are only well defined if ``a`` is guaranteed to be
      # positive definite.
      jtu.check_grads(
          lambda x, y: positive_definite_solve(posify(x), y),
          (a, b),
          order=2,
          rtol=0.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 5 for calling the CustomLinearSolveTest obj's underlying member method code with example usage: obj.test_custom_linear_solve_cholesky() and returns: lax, jsp