code: def main(unused_argv):
      key1, key2 = random.split(random.PRNGKey(1))
  
      x1 = random.normal(key1, (10, 3))
      x2 = random.normal(key2, (20, 3))
  
      # Consider a nonlinearity for which we know the closed-form expression (GeLU).
      _, _, kernel_fn_closed_form = stax.serial(
          stax.Dense(1),
          stax.Gelu(),  # Contains the closed-form GeLU NNGP/NTK expression.
          stax.Dense(1),
      )
      kernel_closed_form = kernel_fn_closed_form(x1, x2)
  
      # Construct the layer from only the elementwise forward-pass GeLU.
      _, _, kernel_fn_numerical = stax.serial(
          stax.Dense(1),
          # Approximation using Gaussian quadrature and autodiff.
          stax.ElementwiseNumerical(jax.nn.gelu, deg=25),
          stax.Dense(1),
      )
      kernel_numerical = kernel_fn_numerical(x1, x2)
  
      # The two kernels are close!
      assert (
          jnp.max(
              jnp.abs(kernel_closed_form.nngp - kernel_numerical.nngp)
          )
          < 1e-3
      )
      assert (
          jnp.max(
              jnp.abs(kernel_closed_form.ntk - kernel_numerical.ntk)
          )
          < 1e-3
      )
      print(
          "Gaussian quadrature approximation of the kernel is accurate!"
      )
  
imports: from absl import app
  from jax import numpy
  from jax import random
  import jax.nn
  from neural_tangents import stax
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) without return types