code: def _concat_masks(
      masks: list[Optional[jnp.ndarray]],
      input_shapes: Sequence[Sequence[int]],
      axis: int,
  ) -> Optional[jnp.ndarray]:
      """Returns a mask which is a concatenation of `masks`.
  
      Since elements of `masks` can have any shapes broadcastable to respective
      elements of `input_shapes`, their concatenation may require broadcasting and
      cannot be done with a single `jnp.concatenate` call.
  
      Args:
        masks: list of masks to concatenate.
        input_shapes: list of input shapes to which the masks are applied.
        axis: concatenation axis.
  
      Returns:
        A single `jnp.ndarray` mask applicable to the concatenated inputs.
      """
      if len(masks) != len(input_shapes):
          raise ValueError(
              f"Number of masks ({len(masks)}) and inputs "
              f"({len(input_shapes)}) don't match, please file a bug at"
              f" https://github.com/google/neural-tangents/issues/new."
          )
  
      if all(m is None for m in masks):
          return None
  
      axis %= len(input_shapes[0])
  
      # Expand the concatenation dimension of each mask.
      masks = [
          m
          if m is None
          else jnp.broadcast_to(
              m,
              (
                  m.shape[:axis]
                  + tuple(input_shapes[i][axis : axis + 1])
                  + m.shape[axis + 1 :]
              ),
          )
          for i, m in enumerate(masks)
      ]
  
      # Max shape to broadcast all masks to along non-concat dimension.
      max_shape = _map_tuples(
          max, (m.shape for m in masks if m is not None)
      )
  
      # Shape of the mask to replace `None` masks with.
      max_shapes = [tuple(map(min, max_shape, i)) for i in input_shapes]
  
      masks = [
          (
              jnp.broadcast_to(
                  m,
                  max_shape[:axis]
                  + m.shape[axis : axis + 1]
                  + max_shape[axis + 1 :],
              )
              if m is not None
              else jnp.zeros_like(max_shapes[i], dtype=jnp.bool_)
          )  # pytype: disable=wrong-arg-types  # jnp-type
          for i, m in enumerate(masks)
      ]
  
      return jnp.concatenate(masks, axis)
  
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 7 for calling a global function with example usage: _concat_masks(masks, input_shapes, axis) and returns: jnp, None