code: def FanInSum() -> InternalLayerMasked:
      """Fan-in sum.
  
      This layer takes a number of inputs (e.g. produced by
      :obj:`~neural_tangents.stax.FanOut`) and sums the inputs to produce a single
      output. Based on :obj:`jax.example_libraries.stax.FanInSum`.
  
      Returns:
        `(init_fn, apply_fn, kernel_fn)`.
      """
      init_fn, apply_fn = ostax.FanInSum
  
      def kernel_fn(ks: Kernels, **kwargs) -> Kernel:
          ks, is_reversed = _preprocess_kernels_for_fan_in(ks)
          if not all(
              [
                  k.shape1 == ks[0].shape1 and k.shape2 == ks[0].shape2
                  for k in ks[1:]
              ]
          ):
              raise ValueError(
                  "All shapes should be equal in `FanInSum/FanInProd`, "
                  f"got `x1.shape`s of {[k.shape1 for k in ks]}, "
                  f"`x2.shape`s of {[k.shape2 for k in ks]}."
              )
  
          is_gaussian = all(k.is_gaussian for k in ks)
          if not is_gaussian and len(ks) != 1:
              # TODO(xlc): FanInSum/FanInConcat could allow non-Gaussian inputs, but
              # we need to propagate the mean of the random variables as well.
              raise NotImplementedError(
                  "`FanInSum` is only implemented for the "
                  "case where all input layers guaranteed to be "
                  "mean-zero Gaussian, i.e. having all "
                  "`is_gaussian` set to `True`, got "
                  f"{[k.is_gaussian for k in ks]}."
              )
  
          _mats_sum = (
              lambda mats: None if mats[0] is None else sum(mats)
          )
  
          cov1s = [k.cov1 for k in ks]
          cov2s = [k.cov2 for k in ks]
          nngps = [k.nngp for k in ks]
          ntks = [k.ntk for k in ks]
          cov1, cov2, nngp, ntk = map(
              _mats_sum, (cov1s, cov2s, nngps, ntks)
          )
  
          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=ks[0].diagonal_batch,
              diagonal_spatial=ks[0].diagonal_spatial,
              shape1=ks[0].shape1,
              shape2=ks[0].shape2,
              batch_axis=ks[0].batch_axis,
              channel_axis=ks[0].channel_axis,
              mask1=None,
              mask2=None,
          )  # pytype:disable=wrong-keyword-args
  
      def mask_fn(mask, input_shape):
          return _sum_masks(mask)
  
      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 2 for calling a global function with example usage: FanInSum() and returns: init_fn,  apply_fn,  kernel_fn,  mask_fn, Kernel, _sum_masks