code: def Elementwise(
      fn: Optional[Callable[[float], float]] = None,
      nngp_fn: Optional[Callable[[float, float, float], float]] = None,
      d_nngp_fn: Optional[
          Callable[[float, float, float], float]
      ] = None,
  ) -> InternalLayer:
      """Elementwise application of `fn` using provided `nngp_fn`.
  
      Constructs a layer given only scalar-valued nonlinearity / activation
      `fn` and the 2D integral `nngp_fn`. NTK function is derived automatically in
      closed form from `nngp_fn`.
  
      If you cannot provide the `nngp_fn`, see :obj:`ElementwiseNumerical` to use
      numerical integration or `nt.monte_carlo.monte_carlo_kernel_fn` to use Monte
      Carlo sampling.
  
      If your function is implemented separately (e.g. `nt.stax.Relu` etc) it's best
      to use the custom implementation, since it uses symbolically simplified
      expressions that are more precise and numerically stable.
  
      For details, please see "`Fast Neural Kernel Embeddings for General
      Activations <https://arxiv.org/abs/2209.04121>`_".
  
      See Also:
        `examples/elementwise.py`.
  
      Example:
        >>> fn = jax.scipy.special.erf  # type: Callable[[float], float]
        >>> #
        >>> def nngp_fn(cov12: float, var1: float, var2: float) -> float:
        >>>   prod = (1 + 2 * var1) * (1 + 2 * var2)
        >>>   return jnp.arcsin(2 * cov12 / np.sqrt(prod)) * 2 / np.pi
        >>> #
        >>> # Use autodiff and vectorization to construct the layer:
        >>> _, _, kernel_fn_auto = stax.Elementwise(fn, nngp_fn)
        >>> #
        >>> # Use custom pre-derived expressions
        >>> # (should be faster and more numerically stable):
        >>> _, _, kernel_fn_stax = stax.Erf()
        >>> #
        >>> kernel_fn_auto(x1, x2) == kernel_fn_stax(x1, x2)  # usually `True`.
  
      Args:
        fn:
          a scalar-input/valued function `fn : R -> R`, the activation /
          nonlinearity. If `None`, invoking the finite width `apply_fn` will raise
          an exception.
  
        nngp_fn:
          a scalar-valued function
          `nngp_fn : (cov12, var1, var2) |-> E[fn(x_1) * fn(x_2)]`, where the
          expectation is over bivariate normal `x1, x2` with variances `var1`,
          `var2` and covarianve `cov12`. Needed for both NNGP and NTK calculation.
          If `None`, invoking infinite width `kernel_fn` will raise an exception.
  
        d_nngp_fn:
          an optional scalar-valued function
          `d_nngp_fn : (cov12, var1, var2) |-> E[fn'(x_1) * fn'(x_2)]` with the same
          `x1, x2` distribution as in `nngp_fn`. If `None`, will be computed using
          automatic differentiation as `d_nngp_fn = d(nngp_fn)/d(cov12)`, which may
          lead to worse precision or numerical stability. `nngp_fn` and `d_nngp_fn`
          are used to derive the closed-form expression for the NTK.
  
      Returns:
        `(init_fn, apply_fn, kernel_fn)`.
  
      Raises:
        NotImplementedError: if a `fn`/`nngp_fn` is not provided, but `apply_fn`/
          `kernel_fn` is called respectively.
      """
      if fn is not None:
          name = fn.__name__
      elif nngp_fn is not None:
          name = nngp_fn.__name__
      else:
          raise ValueError(
              "No finite (`fn`) or infinite (`nngp_fn`) functions "
              "provided, the layer will not do anything."
          )
  
      if nngp_fn is None:
          kernel_fn = None
  
      else:
          if d_nngp_fn is None:
              url = "https://jax.readthedocs.io/en/latest/faq.html#gradients-contain-nan-where-using-where"
              warnings.warn(
                  f"Using JAX autodiff to compute the `fn` derivative for NTK. Beware "
                  f"of {url}."
              )
              d_nngp_fn = jnp.vectorize(grad(nngp_fn))
  
          def kernel_fn(k: Kernel) -> Kernel:
              cov1, nngp, cov2, ntk = k.cov1, k.nngp, k.cov2, k.ntk
  
              var1 = get_diagonal(
                  cov1, k.diagonal_batch, k.diagonal_spatial
              )
              var2 = get_diagonal(
                  cov2, k.diagonal_batch, k.diagonal_spatial
              )
  
              if ntk is not None:
                  ntk *= _vmap_2d(
                      d_nngp_fn,
                      nngp,
                      var1,
                      var2,
                      False,
                      k.diagonal_spatial,
                  )
  
              nngp = _vmap_2d(
                  nngp_fn, nngp, var1, var2, False, k.diagonal_spatial
              )
              cov1 = _vmap_2d(
                  nngp_fn,
                  cov1,
                  var1,
                  None,
                  k.diagonal_batch,
                  k.diagonal_spatial,
              )
              if cov2 is not None:
                  cov2 = _vmap_2d(
                      nngp_fn,
                      cov2,
                      var2,
                      None,
                      k.diagonal_batch,
                      k.diagonal_spatial,
                  )
              return k.replace(cov1=cov1, nngp=nngp, cov2=cov2, ntk=ntk)
  
      return _elementwise(fn, name, kernel_fn)
  
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 20 for calling a global function with example usage: Elementwise(fn, nngp_fn, d_nngp_fn) and returns: _elementwise, k