code: def requires(**static_reqs):
      """Returns a decorator that augments `kernel_fn` with consistency checks.
  
      Use this to specify your `kernel_fn` input kernel requirements.
  
      See Also:
        :class:`Diagonal`, :class:`Bool`.
  
      """
  
      def req(kernel_fn: LayerKernelFn):
          """Returns `kernel_fn` with additional consistency checks."""
  
          @utils.wraps(kernel_fn)
          def new_kernel_fn(
              k: NTTree[Kernel], **kwargs
          ) -> NTTree[Kernel]:
              """Executes `kernel_fn` on `kernels` after checking consistency."""
              fused_reqs = _fuse_requirements(static_reqs, {}, **kwargs)
  
              # `FanInConcat / FanInSum` have no requirements and
              # execute custom consistency checks.
              if isinstance(k, Kernel):
                  for key, v in fused_reqs.items():
                      if (
                          v is not None
                      ):  # `None` is treated as explicitly not having a req.
                          if key in (
                              "diagonal_batch",
                              "diagonal_spatial",
                          ):
                              if getattr(k, key) is True and (
                                  v is False
                                  or (
                                      isinstance(v, Diagonal)
                                      and v.input == Bool.NO
                                  )
                              ):
                                  raise ValueError(
                                      f"{kernel_fn} requires `{key} == {v}`, but "
                                      f"input kernel has `{key} == True`, hence "
                                      f"does not contain sufficient information. "
                                      f"Please recompute the input kernel with "
                                      f"`{key} == {v}`."
                                  )
  
                          elif key in ("batch_axis", "channel_axis"):
                              ndim = len(
                                  k.shape1
                              )  # pytype: disable=attribute-error  # preserve-union-macros
                              v_kernel = getattr(k, key)
                              v_pos = v % ndim
                              if v_kernel != v_pos:
                                  raise ValueError(
                                      f"{kernel_fn} requires `{key} == {v_pos}`, "
                                      f"but input kernel has `{key} == {v_kernel}`, "
                                      f"making the infinite limit ill-defined."
                                  )
  
                          else:
                              # Any other name is recognized as a keyword-argument threaded
                              # through all `kernel_fn` down to `_inputs_to_kernel` rather than
                              # a requirement for this layer.
                              pass
  
              return kernel_fn(k, **kwargs)
  
          _set_req(new_kernel_fn, frozendict.frozendict(static_reqs))
          return new_kernel_fn
  
      return req
  
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 2 for calling a global function with example usage: requires() and returns: req, new_kernel_fn, kernel_fn