code: def ElementwiseNumerical(
      fn: Callable[[float], float],
      deg: int,
      df: Optional[Callable[[float], float]] = None,
  ) -> InternalLayer:
      """Activation function using numerical integration.
  
      Supports general activation functions using Gauss-Hermite quadrature.
  
      For details, please see "`Fast Neural Kernel Embeddings for General
      Activations <https://arxiv.org/abs/2209.04121>`_".
  
      See Also:
        `examples/elementwise_numerical.py`.
  
      Args:
        fn:
          activation function.
  
        deg:
          number of sample points and weights for quadrature. It must be >= 1.
          We observe for smooth activations `deg=25` is a good place to start.
          For non-smooth activation functions (e.g. ReLU, Abs) quadrature is not
          recommended (for now use `nt.monte_carlo_kernel_fn`). Due to bivariate
          integration, compute time and memory scale as O(deg**2) for more
          precision. See eq (13) in
          https://mathworld.wolfram.com/Hermite-GaussQuadrature.html
          for error estimates in the case of 1d Gauss-Hermite quadrature.
  
        df:
          optional, derivative of the activation function (`fn`). If not provided,
          it is computed by `jax.grad`. Providing analytic derivative can speed up
          the NTK computations.
  
      Returns:
        `(init_fn, apply_fn, kernel_fn)`.
      """
      warnings.warn(
          f"Numerical Activation Layer with fn={fn}, deg={deg} used!"
          "Note that numerical error is controlled by `deg` and for a given"
          "tolerance level, required `deg` will highly be dependent on the choice"
          "of `fn`."
      )
  
      quad_points = sp.special.roots_hermite(deg)
  
      if df 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 of "
              f"{url}."
          )
          df = jnp.vectorize(grad(fn))
  
      def kernel_fn(k: Kernel) -> Kernel:
          """Kernel transformation of activation function using quadrature."""
          cov1, nngp, cov2, ntk = k.cov1, k.nngp, k.cov2, k.ntk
  
          d1 = get_diagonal(cov1, k.diagonal_batch, k.diagonal_spatial)
          d2 = get_diagonal(cov2, k.diagonal_batch, k.diagonal_spatial)
  
          end_axis = 1 if k.diagonal_spatial else cov1.ndim
          q11 = utils.interleave_ones(d1, 0, end_axis, True)
          q22 = utils.interleave_ones(
              d1 if d2 is None else d2, 0, end_axis, False
          )
  
          def nngp_ntk_fn(nngp, q11, q22, ntk=None):
              """Simple Gauss-Hermite quadrature routine."""
              xs, ws = quad_points
              grid = jnp.outer(ws, ws)
              x = xs.reshape((xs.shape[0],) + (1,) * (nngp.ndim + 1))
              y = xs.reshape((1, xs.shape[0]) + (1,) * nngp.ndim)
              xy_axes = (0, 1)
  
              nngp = jnp.expand_dims(nngp, xy_axes)
              q11, q22 = jnp.expand_dims(q11, xy_axes), jnp.expand_dims(
                  q22, xy_axes
              )
  
              def integrate(f):
                  fvals = f(
                      _sqrt(2 * q11) * x
                  ) * f(  # pytype: disable=wrong-arg-types  # jnp-type
                      nngp / _sqrt(q11 / 2, 1e-30) * x
                      + _sqrt(2 * (q22 - nngp**2 / q11)) * y
                  )
                  return (
                      jnp.tensordot(grid, fvals, (xy_axes, xy_axes))
                      / jnp.pi
                  )
  
              if ntk is not None:
                  ntk *= integrate(df)
              nngp = integrate(fn)
              return nngp, ntk
  
          def nngp_fn_diag(nngp):
              xs, ws = quad_points
              x = xs.reshape((xs.shape[0],) + (1,) * nngp.ndim)
              x_axes = (0,)
              nngp = jnp.expand_dims(nngp, x_axes)
              fval = fn(_sqrt(2 * nngp) * x) ** 2
              return jnp.tensordot(
                  ws, fval, (x_axes, x_axes)
              ) / jnp.sqrt(jnp.pi)
  
          nngp, ntk = nngp_ntk_fn(nngp, q11, q22, ntk)
  
          if k.diagonal_batch and k.diagonal_spatial:
              cov1 = nngp_fn_diag(cov1)
              if cov2 is not None:
                  cov2 = nngp_fn_diag(cov2)
  
          else:
              start_axis = 1 if k.diagonal_batch else 0
              q11 = utils.interleave_ones(
                  d1, start_axis, end_axis, True
              )
              q22 = utils.interleave_ones(
                  d1, start_axis, end_axis, False
              )
              cov1, _ = nngp_ntk_fn(cov1, q11, q22)
  
              if cov2 is not None:
                  q11 = utils.interleave_ones(
                      d2, start_axis, end_axis, True
                  )
                  q22 = utils.interleave_ones(
                      d2, start_axis, end_axis, False
                  )
                  cov2, _ = nngp_ntk_fn(cov2, q11, q22)
  
          return k.replace(cov1=cov1, nngp=nngp, cov2=cov2, ntk=ntk)
  
      return _elementwise(
          fn, f"ElementwiseNumerical({fn},deg={deg})", 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 21 for calling a global function with example usage: ElementwiseNumerical(fn, deg, df) and returns: _elementwise, k, nngp,  ntk, unknown, unknown