code: def _get_cho_solve(
      A: jnp.ndarray,
      diag_reg: float,
      diag_reg_absolute_scale: bool,
      lower: bool = False,
  ) -> Callable[[jnp.ndarray, Axes], jnp.ndarray]:
      x_non_channel_shape = A.shape[1::2]
      A = utils.make_2d(A)
      A = _add_diagonal_regularizer(
          A, diag_reg, diag_reg_absolute_scale
      )
      C = jsp.linalg.cho_factor(A, lower)
  
      def cho_solve(b: jnp.ndarray, b_axes: Axes) -> jnp.ndarray:
          b_axes = utils.canonicalize_axis(b_axes, b)
          last_b_axes = range(-len(b_axes), 0)
          x_shape = x_non_channel_shape + tuple(
              b.shape[a] for a in b_axes
          )
  
          b = jnp.moveaxis(b, b_axes, last_b_axes)
          b = b.reshape((A.shape[1], -1))
  
          x = jsp.linalg.cho_solve(C, b)
          x = x.reshape(x_shape)
          return x
  
      return cho_solve
  
imports: import collections
  from functools import lru_cache
  from typing import Any, Callable, Generator, Iterable, NamedTuple, Optional, Protocol, Union
  import jax
  from jax import grad
  from jax.experimental import ode
  import jax.numpy
  import jax.scipy
  from jax.tree_util import tree_all
  from jax.tree_util import tree_map
  import numpy
  import scipy
  from .utils import dataclasses
  from .utils import utils
  from .utils.typing import Axes
  from .utils.typing import Get
  from .utils.typing import KernelFn
label_desc: Use image node_id 10 for calling a global function with example usage: _get_cho_solve(A, diag_reg, diag_reg_absolute_scale, lower) and returns: cho_solve, x