code: def FanInConcat(axis: int = -1) -> InternalLayerMasked:
      """Fan-in concatenation.
  
      This layer takes a number of inputs (e.g. produced by
      :obj:`~neural_tangents.stax.FanOut`) and concatenates the inputs to produce a
      single output. Based on :obj:`jax.example_libraries.stax.FanInConcat`.
  
      Args:
        axis: Specifies the axis along which input tensors should be concatenated.
  
      Returns:
        `(init_fn, apply_fn, kernel_fn)`.
      """
      init_fn, apply_fn = ostax.FanInConcat(axis)
  
      def kernel_fn(ks: Kernels, **kwargs) -> Kernel:
          ks, is_reversed = _preprocess_kernels_for_fan_in(ks)
  
          diagonal_batch = ks[0].diagonal_batch
          diagonal_spatial = ks[0].diagonal_spatial
  
          shape1, shape2 = ks[0].shape1, ks[0].shape2
  
          ndim = len(shape1)
          _axis = axis % ndim
          batch_axis = ks[0].batch_axis
          channel_axis = ks[0].channel_axis
  
          new_shape1 = shape1[:_axis] + shape1[_axis + 1 :]
          new_shape2 = shape2[:_axis] + shape2[_axis + 1 :]
          for k in ks:
              k_shape1 = k.shape1[:_axis] + k.shape1[_axis + 1 :]
              k_shape2 = k.shape2[:_axis] + k.shape2[_axis + 1 :]
              if k_shape1 != new_shape1 or k_shape2 != new_shape2:
                  raise ValueError(
                      "Non-`axis` shapes should be equal in `FanInConcat`."
                  )
  
          # Check if inputs are independent Gaussians.
          if _axis != channel_axis:
              is_gaussian = all(k.is_gaussian for k in ks)
              if not is_gaussian:
                  # TODO(xlc): FanInSum/FanInConcat could allow non-Gaussian inputs, but
                  # we need to propagate the mean of the random variables as well.
                  raise NotImplementedError(
                      "`FanInConcat` layer along the non-channel axis is only implemented"
                      "for the case if all input layers guaranteed to be mean-zero "
                      "Gaussian, i.e. having all `is_gaussian` set to `True`."
                  )
          else:
              # TODO(romann): allow nonlinearity after channelwise concatenation.
              # TODO(romann): support concatenating different channelwise masks.
              is_gaussian = False
  
          if _axis == batch_axis:
              warnings.warn(
                  f"Concatenation along the batch axis ({_axis}) gives "
                  f"inconsistent covariances when batching - "
                  f"proceed with caution."
              )
  
          spatial_axes = tuple(
              i
              for i in range(ndim)
              if i not in (channel_axis, batch_axis)
          )
          # Change spatial axis according to the kernel `is_reversed`.
          if _axis in spatial_axes and is_reversed:
              _axis = spatial_axes[::-1][spatial_axes.index(_axis)]
  
          # Map activation tensor axis to the covariance tensor axis.
          tensor_axis_to_kernel_axis = {
              **{
                  batch_axis: 0,
                  channel_axis: -1,
              },
              **{
                  spatial_axis: idx + 1
                  for idx, spatial_axis in enumerate(spatial_axes)
              },
          }
  
          _axis = tensor_axis_to_kernel_axis[_axis]
          widths = [k.shape1[channel_axis] for k in ks]
  
          cov1 = _concat_kernels(
              [k.cov1 for k in ks],
              _axis,
              diagonal_batch,
              diagonal_spatial,
              widths,
          )
          cov2 = _concat_kernels(
              [k.cov2 for k in ks],
              _axis,
              diagonal_batch,
              diagonal_spatial,
              widths,
          )
          nngp = _concat_kernels(
              [k.nngp for k in ks],
              _axis,
              False,
              diagonal_spatial,
              widths,
          )
          ntk = _concat_kernels(
              [k.ntk for k in ks],
              _axis,
              False,
              diagonal_spatial,
              widths,
          )
  
          return Kernel(
              cov1=cov1,
              cov2=cov2,
              nngp=nngp,
              ntk=ntk,
              x1_is_x2=ks[0].x1_is_x2,
              is_gaussian=is_gaussian,
              is_reversed=is_reversed,
              is_input=ks[0].is_input,
              diagonal_batch=diagonal_batch,
              diagonal_spatial=diagonal_spatial,
              shape1=None,
              shape2=None,
              batch_axis=batch_axis,
              channel_axis=channel_axis,
              mask1=None,
              mask2=None,
          )  # pytype:disable=wrong-keyword-args
  
      def mask_fn(mask, input_shape):
          return _concat_masks(mask, input_shape, axis)
  
      return init_fn, apply_fn, kernel_fn, mask_fn
  
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 4 for calling a global function with example usage: FanInConcat(axis) and returns: init_fn,  apply_fn,  kernel_fn,  mask_fn, Kernel, _concat_masks