code: def _concat_kernels(
      mats: Sequence[Optional[jnp.ndarray]],
      axis: int,
      diagonal_batch: bool,
      diagonal_spatial: bool,
      widths: Sequence[int],
  ) -> Optional[jnp.ndarray]:
      """Compute the covariance of concatenated activations with given covariances.
  
      Args:
        mats: Covariance tensors of the same shape.
  
        axis: Specifies the axis along which the covariances (not activations) are
          concatenated. `-1` corresponds to averaging.
  
        diagonal_batch: Specifies whether `cov1` and `cov2` store only
          the diagonal of the sample-sample covariance
          (`diagonal_batch == True`,
           `cov1.shape == (batch_size_1, ...)`),
          or the full covariance
          (`diagonal_batch == False`,
           `cov1.shape == (batch_size_1, batch_size_1, ...)`).
  
        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, ...)`).
  
        widths: list of integer channel widths of the finite model inputs.
  
      Returns:
        New `jnp.ndarray` representing covariance between concatenated activations.
      """
      if mats[0] is None:
          return None
  
      n_mats = len(mats)
      mat_ndim = mats[0].ndim
  
      # Averaging if concatenating along features or diagonalized dimension.
      if axis == -1:
          if all(w == widths[0] for w in widths):
              widths = [1] * len(widths)
          mat = sum(mats[i] * widths[i] for i in range(n_mats)) / sum(
              widths
          )
  
      # Simple concatenation along the axis if the axis is not duplicated.
      elif (axis == 0 and diagonal_batch) or (
          axis != 0 and diagonal_spatial
      ):
          concat_axis = axis + (0 if diagonal_batch else 1)
          mat = jnp.concatenate(mats, concat_axis)
  
      # 2D concatenation with insertion of 0-blocks if the axis is present twice.
      else:
          rows = []
          pad_axis = max(0, 2 * axis - (1 if diagonal_batch else 0))
          for i, mat in enumerate(mats):
              pads = [(0, 0)] * mat_ndim
              pads[pad_axis] = (
                  sum(mats[j].shape[pad_axis] for j in range(i)),
                  sum(
                      mats[j].shape[pad_axis]
                      for j in range(i + 1, n_mats)
                  ),
              )
              rows.append(jnp.pad(mat, pads))
          mat = jnp.concatenate(rows, pad_axis + 1)
  
      return mat
  
imports: import functools
  from typing import Callable, Iterable, Optional, Sequence
  import warnings
  from jax import numpy
  import jax.example_libraries.stax
  from ..utils.kernel import Kernel
  from ..utils.typing import InternalLayer
  from ..utils.typing import InternalLayerMasked
  from ..utils.typing import Kernels
  from .requirements import layer
  from .requirements import supports_masking
label_desc: Use image node_id 9 for calling a global function with example usage: _concat_kernels(mats, axis, diagonal_batch, diagonal_spatial, widths) and returns: mat, None