code: def _vmap_2d(
      fn: Callable[[float, float, float], float],
      cov12: jnp.ndarray,
      var1: jnp.ndarray,
      var2: Optional[jnp.ndarray],
      diagonal_batch: bool,
      diagonal_spatial: bool,
  ) -> jnp.ndarray:
      """Effectively a "2D vmap" of `fn(cov12, var1, var2)`.
  
      Applicable for all possible kernel layouts.
  
      Args:
        fn:
          scalar-valued, elementwise `fn(cov12, var1, var2)` function to apply.
  
        cov12:
          covariance tensor (`q12`), `nngp`/`ntk`/`cov1`/`cov2`, of shape
          `(N1[, N2])`, `(N1[, N2], X, Y, ...)`, `(N1[, N2], X, X, Y, Y, ...)`
          depending on `diagonal_batch`, `diagonal_spatial`, and the number of
          spatial dimensions.
  
        var1:
          variance tensor (`q11`), has shape `(N1[, X, Y, ...])`.
  
        var2:
          variance tensor (`q22`), has shape `(N1[, X, Y, ...])`.
  
        diagonal_batch:
          `True` if `cov12` has only one batch dimension.
  
        diagonal_spatial:
          `True` if `cov12` has spatial dimensions appearing once (vs twice).
  
      Returns:
        Resulting array `[fn(cov12[i, j], var1[i], var2[j])]_{i j}`. Has the same
        shape as `cov12`.
      """
      batch_ndim = 1 if diagonal_batch else 2
      start = 2 - batch_ndim
      cov_end = batch_ndim if diagonal_spatial else cov12.ndim
      _cov12 = utils.make_2d(cov12, start, cov_end)
  
      var_end = 1 if diagonal_spatial else var1.ndim
      var1 = var1.reshape(
          var1.shape[:start] + (-1,) + var1.shape[var_end:]
      )
      var2 = (
          var1
          if var2 is None
          else var2.reshape(
              var2.shape[:start] + (-1,) + var2.shape[var_end:]
          )
      )
  
      fn = vmap(
          vmap(
              jnp.vectorize(fn),
              in_axes=(start, None, start),
              out_axes=start,
          ),
          in_axes=(start, start, None),
          out_axes=start,
      )
      out = fn(_cov12, var1, var2)  # type: jnp.ndarray
      out_shape = (
          cov12.shape[:start]
          + cov12.shape[start:cov_end:2]
          + cov12.shape[start + 1 : cov_end : 2]
          + cov12.shape[cov_end:]
      )
      out = out.reshape(out_shape)
      out = utils.zip_axes(out, start, cov_end)
      return out
  
imports: import functools
  import operator
  from typing import Callable, Optional, Sequence
  import warnings
  import jax
  from jax import custom_jvp
  from jax import grad
  from jax import numpy
  from jax import vmap
  from jax.scipy.special import erf
  import numpy
  import scipy
  from ..utils import utils
  from ..utils.kernel import Kernel
  from ..utils.typing import InternalLayer
  from ..utils.typing import LayerKernelFn
  from .requirements import Diagonal
  from .requirements import get_diagonal
  from .requirements import get_diagonal_outer_prods
  from .requirements import layer
  from .requirements import requires
  from .requirements import supports_masking
label_desc: Use image node_id 27 for calling a global function with example usage: _vmap_2d(fn, cov12, var1, var2, diagonal_batch, diagonal_spatial) and returns: out