code: def _inputs_to_kernel(
      x1: jnp.ndarray,
      x2: Optional[jnp.ndarray],
      *,
      diagonal_batch: bool,
      diagonal_spatial: Union[bool, Diagonal],
      compute_ntk: bool,
      batch_axis: int,
      channel_axis: Optional[int],
      mask_constant: Optional[float],
      eps: float = 1e-12,
      **kwargs,
  ) -> Kernel:
      """Transforms (batches of) inputs to a `Kernel`.
  
      This is a private function. Docstring and example are for internal reference.
  
      The kernel contains the empirical covariances between different inputs and
      their entries (e.g. pixels, words, entries in a time series etc.) necessary
      to compute the covariance of the Gaussian Process corresponding to an
      infinite Bayesian or continuous gradient descent trained neural network.
  
      The smallest necessary number of covariance entries is tracked. For example,
      all networks are assumed to have i.i.d. weights along the channel / feature
      / logits dimensions, hence covariance between different entries along these
      dimensions is known to be 0 and is not tracked.
  
      Example:
        >>> x = jnp.ones((10, 32, 16, 3))
        >>> o = _inputs_to_kernel(x, None,
        >>>                      diagonal_batch=True,
        >>>                      diagonal_spatial=False,
        >>>                      compute_ntk=True,
        >>>                      batch_axis=0,
        >>>                      channel_axis=-1)
        >>> o.cov1.shape, o.ntk.shape
        (10, 32, 32, 16, 16), (10, 10, 32, 32, 16, 16)
        >>> o = _inputs_to_kernel(x, None,
        >>>                      diagonal_batch=True,
        >>>                      diagonal_spatial=True,
        >>>                      compute_ntk=True,
        >>>                      batch_axis=0,
        >>>                      channel_axis=-1)
        >>> o.cov1.shape, o.ntk.shape
        (10, 32, 16), (10, 10, 32, 16)
        >>> x1 = jnp.ones((10, 128))
        >>> x2 = jnp.ones((20, 128))
        >>> o = _inputs_to_kernel(x1, x2,
        >>>                      diagonal_batch=True,
        >>>                      diagonal_spatial=True,
        >>>                      compute_ntk=False,
        >>>                      batch_axis=0,
        >>>                      channel_axis=-1)
        >>> o.cov1.shape, o.nngp.shape
        (10,), (10, 20)
  
      Args:
        x1:
          an `(S+2)`-dimensional `jnp.ndarray` of shape
          `(batch_size_1, height, width, depth, ..., n_channels)` with `S` spatial
          dimensions (`S >= 0`). Dimensions may be in different order based on
          `batch_axis` and `channel_axis`.
  
        x2:
          an optional `jnp.ndarray` with the same shape as `x1` apart from possibly
          different batch size. `None` means `x2 == x1`.
  
        diagonal_batch:
          Specifies whether `cov1` and `cov2` 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, ...)`).
  
        diagonal_spatial:
          Specifies whether all (`cov1`, `ntk`, etc.) input covariance matrcies
          should store only the diagonals of the location-location covariances
          (`diagonal_spatial == True`,
           `nngp.shape == (batch_size_1, batch_size_2, height, width, depth, ...)`),
          or the full covariance
          (`diagonal_spatial == False`,
           `nngp.shape == (batch_size_1, batch_size_2, height, height,
                           width, width, depth, depth, ...)`).
  
        compute_ntk:
          `True` to compute both NTK and NNGP kernels, `False` to only compute NNGP.
  
        batch_axis:
          Specifies which axis is the batch axis.
  
        channel_axis:
          Specifies which axis is the channel / feature axis. For `kernel_fn`,
          channel size is considered to be infinite.
  
        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.
  
        eps:
          a small number used to check whether x1 and x2 are the same up to `eps`.
  
        **kwargs:
          other arguments passed to all intermediary `kernel_fn` calls (not used
          here).
  
      Returns:
        The :class:`~neural_tangents.Kernel` object containing inputs covariance[s].
      """
  
      if not (
          isinstance(x1, (np.ndarray, jnp.ndarray))
          and (x2 is None or isinstance(x2, (np.ndarray, jnp.ndarray)))
      ):
          raise TypeError(
              (
                  "Wrong input types given. Found `x1` of type "
                  f"{type(x1)} and `x2` of type {type(x2)}, need both to be"
                  f"`jnp.ndarray`s (`x2` can be `None`)."
              )
          )
  
      batch_axis %= x1.ndim
      diagonal_spatial = bool(diagonal_spatial)
  
      if batch_axis != 0:
          # TODO(romann): add support or clear error for batching.
          warnings.warn(
              f"!!! Non-leading (!= 0) batch dimension in the "
              f"input layer is not supported for batching "
              f"kernels, got batch_axis = {batch_axis}. !!!"
          )
  
      if channel_axis is None:
  
          def flatten(x):
              if x is None:
                  return x
              return jnp.moveaxis(x, batch_axis, 0).reshape(
                  (x.shape[batch_axis], -1)
              )
  
          x1, x2 = flatten(x1), flatten(x2)
          batch_axis, channel_axis = 0, 1
          diagonal_spatial = False
  
      else:
          channel_axis %= x1.ndim
  
      def get_x_cov_mask(x):
          if x is None:
              return None, None, None
  
          if x.ndim < 2:
              raise ValueError(
                  f"Inputs must be at least 2D (a batch dimension and a "
                  f"channel/feature dimension), got {x.ndim}."
              )
  
          x = _get_masked_array(x, mask_constant)
          x, mask = x.masked_value, x.mask
  
          # TODO(schsam): Think more about dtype automatic vs manual dtype promotion.
          x = x.astype(jax.dtypes.canonicalize_dtype(jnp.float64))
  
          if diagonal_batch:
              cov = _cov_diag_batch(
                  x, diagonal_spatial, batch_axis, channel_axis
              )
          else:
              cov = _cov(
                  x, x, diagonal_spatial, batch_axis, channel_axis
              )
  
          return x, cov, mask
  
      x1, cov1, mask1 = get_x_cov_mask(x1)
      x2, cov2, mask2 = get_x_cov_mask(x2)
      nngp = _cov(x1, x2, diagonal_spatial, batch_axis, channel_axis)
  
      ntk = (
          jnp.zeros((), nngp.dtype) if compute_ntk else None
      )  # pytype: disable=attribute-error  # always-use-return-annotations
      is_gaussian = False
      is_reversed = False
      x1_is_x2 = utils.x1_is_x2(x1, x2, eps=eps)
      is_input = False
  
      return Kernel(
          cov1=cov1,
          cov2=cov2,
          nngp=nngp,
          ntk=ntk,
          x1_is_x2=x1_is_x2,
          is_gaussian=is_gaussian,
          is_reversed=is_reversed,
          is_input=is_input,
          diagonal_batch=diagonal_batch,
          diagonal_spatial=diagonal_spatial,
          shape1=x1.shape,
          shape2=x1.shape if x2 is None else x2.shape,
          batch_axis=batch_axis,
          channel_axis=channel_axis,
          mask1=mask1,
          mask2=mask2,
      )
  
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 15 for calling a global function with example usage: _inputs_to_kernel(x1, x2) and returns: Kernel, x,  cov,  mask, jnp, None,  None,  None, x