code: def _preprocess_kernel_fn(
      init_fn: InitFn, apply_fn: ApplyFn, kernel_fn: LayerKernelFn
  ) -> AnalyticKernelFn:
      """Returns a `kernel_fn` with additional arguments.
  
      Args:
        init_fn: layer parameters initialization function. Used for shape inference.
  
        apply_fn: layer forward-prop function. Used for shape inference.
  
        kernel_fn: the `Kernel` -> `Kernel` layer propagation function.
  
      Returns:
        A new `kernel_fn` that does the same computation but accepts additional
        arguments to flexibly specify the required computation, and can be applied
        to either a `Kernel' or a pair of `jnp.ndarrray`s.
      """
      # Set empty requirements if none specified.
      if not _has_req(kernel_fn):
          kernel_fn = requires()(kernel_fn)
  
      def kernel_fn_kernel(kernel, **kwargs):
          out_kernel = kernel_fn(kernel, **kwargs)
          return _set_shapes(
              init_fn, apply_fn, kernel, out_kernel, **kwargs
          )
  
      def kernel_fn_x1(x1, x2, get, **kwargs):
          # Get input requirements requested by network layers, user, or defaults.
          kernel_fn_reqs = get_req(kernel_fn)
          reqs = _fuse_requirements(
              kernel_fn_reqs, _DEFAULT_INPUT_REQ, **kwargs
          )
          compute_ntk = (get is None) or ("ntk" in get)
  
          if x2 is None:
              x2 = tree_map(lambda x: None, x1)
  
          def input_fn(x1, x2):
              return _inputs_to_kernel(
                  x1, x2, compute_ntk=compute_ntk, **reqs
              )
  
          kernel = tree_map(input_fn, x1, x2)
  
          out_kernel = kernel_fn(kernel, **kwargs)
          return _set_shapes(
              init_fn, apply_fn, kernel, out_kernel, **kwargs
          )
  
      @utils.get_namedtuple("AnalyticKernel")
      def kernel_fn_any(
          x1_or_kernel: Union[NTTree[jnp.ndarray], NTTree[Kernel]],
          x2: Optional[NTTree[jnp.ndarray]] = None,
          get: Optional[Get] = None,
          *,
          pattern: Optional[
              tuple[Optional[jnp.ndarray], Optional[jnp.ndarray]]
          ] = None,
          mask_constant: Optional[float] = None,
          diagonal_batch: Optional[bool] = None,
          diagonal_spatial: Optional[bool] = None,
          **kwargs,
      ):
          """Returns the `Kernel` resulting from applying `kernel_fn` to given inputs.
  
          Args:
            x1_or_kernel:
              either an NTTree of the first batch of inputs.
  
            x2:
              an optional NTTree of `jnp.ndarray` with the second batch of inputs.
              `None` means `x2 == x1` or `x1_or_kernel is Kernel`.
  
            get:
              either `None`, a string, or a tuple of strings specifying which data
              should be returned by the kernel function. Can be "nngp", "ntk", "cov1",
              "cov2", "is_gaussian", "is_reversed", "diagonal_batch",
              "diagonal_spatial", etc.
  
            pattern:
              either `None` or a tuple of two `jnp.ndarray`. The
              `pattern = (pattern1, pattern2)` is used to specify how the nodes in a
              graphical network is aggregated.
  
            mask_constant:
              an optional `float`, the value in inputs to be considered
              as masked (e.g. padding in a batch of sentences). `None` means no
              masking. Can also be `jnp.nan`, `jnp.inf` etc. Beware of floating point
              precision errors and try to use an atypical for inputs value.
  
            diagonal_batch:
              an optional boolean specifying whether `cov1` and `cov2` in all
              intermediary layers should 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, ...)`).
              Defaults to least compute-heavy setting necessary to compute the output
              `nngp` [and `ntk`] covariance.
  
            diagonal_spatial:
              an optional boolean specifying whether all (`cov1`, `ntk`, etc.)
              covariance matrcies in all intermediary layers should store only the
              diagonals of the location-location covariances
              (`diagonal_spatial == True`,
               `nngp.shape == (batch_size_1, batch_size_2, height, width, ...)`),
              or the full covariance
              (`diagonal_spatial == False`,
               `nngp.shape == (batch_size_1, batch_size_2, height, height,
                               width, width, ...)`).
              Defaults to least compute-heavy setting necessary to compute the output
              `nngp` [and `ntk`] covariance.
  
            **kwargs:
              other arguments passed to all intermediary `kernel_fn` calls.
  
          Returns:
            If `get` is a string, returns the requested `jnp.ndarray`. If `get` is a
            tuple, returns an `AnalyticKernel` namedtuple containing only the
            requested information. If `get` is `None` then a `Kernel` object is
            returned containing all the data.
          """
  
          def all_of(x, cls: type) -> bool:
              def is_leaf(x) -> bool:
                  return isinstance(
                      x, (Kernel, jnp.ndarray, np.ndarray)
                  )
  
              return tree_all(
                  tree_map(
                      lambda x: isinstance(x, cls), x, is_leaf=is_leaf
                  )
              )
  
          if all_of(x1_or_kernel, Kernel) and x2 is None:
              return kernel_fn_kernel(
                  x1_or_kernel,
                  pattern=pattern,
                  diagonal_batch=diagonal_batch,
                  diagonal_spatial=diagonal_spatial,
                  **kwargs,
              )
  
          return kernel_fn_x1(
              x1_or_kernel,
              x2,
              get,
              pattern=pattern,
              diagonal_batch=diagonal_batch,
              diagonal_spatial=diagonal_spatial,
              mask_constant=mask_constant,
              **kwargs,
          )
  
      _set_req(kernel_fn_any, get_req(kernel_fn))
      return kernel_fn_any
  
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 19 for calling a global function with example usage: _preprocess_kernel_fn(init_fn, apply_fn, kernel_fn) and returns: kernel_fn_any, _set_shapes, _set_shapes, kernel_fn_x1, _inputs_to_kernel, tree_all, kernel_fn_kernel, isinstance