code: def get_diagonal(
      cov: Optional[jnp.ndarray],
      diagonal_batch: bool,
      diagonal_spatial: bool,
  ) -> Optional[jnp.ndarray]:
      """Extracts the diagonal of `cov` over all (sample, spatial) dimensions.
  
      Adapts computation if `cov` already stores only the diagonal along some
      dimensions based on `diagonal_batch` and `diagonal_spatial`.
      """
      if cov is None:
          return cov
  
      batch_ndim = 1 if diagonal_batch else 2
      start_axis = 2 - batch_ndim
      end_axis = batch_ndim if diagonal_spatial else cov.ndim
      cov = utils.unzip_axes(cov, start_axis, end_axis)
      return utils.diagonal_between(cov, start_axis, end_axis)
  
imports: import dataclasses
  import enum
  from typing import Callable, Optional, Sequence, Union
  import warnings
  import frozendict
  import jax
  from jax import eval_shape
  from jax import lax
  from jax import numpy
  from jax.core import ShapedArray
  from jax.tree_util import tree_all
  from jax.tree_util import tree_map
  import numpy
  from ..utils import dataclasses
  from ..utils import utils
  from ..utils.kernel import Kernel
  from ..utils.typing import AnalyticKernelFn
  from ..utils.typing import ApplyFn
  from ..utils.typing import Axes
  from ..utils.typing import Get
  from ..utils.typing import InitFn
  from ..utils.typing import InternalLayer
  from ..utils.typing import Layer
  from ..utils.typing import LayerKernelFn
  from ..utils.typing import NTTree
  from ..utils.typing import PyTree
label_desc: Use image node_id 20 for calling a global function with example usage: get_diagonal(cov, diagonal_batch, diagonal_spatial) and returns: utils, cov