code: def _cov(
      x1: jnp.ndarray,
      x2: Optional[jnp.ndarray],
      diagonal_spatial: bool,
      batch_axis: int,
      channel_axis: int,
  ) -> Optional[jnp.ndarray]:
      """Computes uncentered covariance (nngp) between two batches of inputs.
  
      Args:
        x1:
          a (2+S)D (S >= 0) `jnp.ndarray` of shape
          `(batch_size_1, <S spatial dimensions>, n_channels)`. `batch_size_1`,
          `n_channels` may be in different positions based on `batch_axis` and
          `channel_axis`.
  
        x2:
          an optional `jnp.ndarray` that has the same shape as `a` apart from
          possibly different batch (`batch_size_2`) dimension. `None` means
          `x2 == x1`.
  
        diagonal_spatial:
          Specifies whether only the diagonals of the
          location-location covariances will be computed,
          (`diagonal_spatial == True`,
           `nngp.shape == (batch_size_1, batch_size_2, height, width, depth, ...)`),
          or the full covariance
          (`diagonal_spatial == False`,
           `nngp.shape == (batch_size_1, batch_size_2, height, height,
                           width, width, depth, depth, ...)`).
  
        batch_axis:
          Specifies which axis is the batch axis.
  
        channel_axis:
          Specifies which axis is the channel / feature axis. For `kernel_fn`,
          channel size is considered to be infinite.
  
      Returns:
        Matrix of uncentred batch covariances with shape
        `(batch_size_1, batch_size_2, <S spatial dimensions>)`
        if `diagonal_spatial` is `True`, or
        `(batch_size_1, batch_size_2, <2*S spatial dimensions>)`
        if `diagonal_spatial` is `False`.
      """
      x2 = x1 if x2 is None else x2
  
      if diagonal_spatial:
          ret = _cov_full_batch_diag_spatial(
              x1, x2, batch_axis, channel_axis
          )
  
      else:
          ret = _cov_full_batch_full_spatial(
              x1, x2, batch_axis, channel_axis
          )
  
      return ret / x1.shape[channel_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 14 for calling a global function with example usage: _cov(x1, x2, diagonal_spatial, batch_axis, channel_axis) and returns: unknown