code: def main(unused_argv):
      # Consider the normalized exponential kernel from
      # https://arxiv.org/abs/2003.02237 (page 6).
      def nngp_fn(cov12, var1, var2):
          prod = jnp.sqrt(var1 * var2)
          return prod * jnp.exp(cov12 / prod - 1)
  
      # This kernel has no known corresponding elementwise nonlinearity.
      # `stax.Elementwise` derives the NTK kernel automatically under the hood using
      # automatic differentiation, without the need to know the respective
      # nonlinearity or computing the integrals by hand.
      _, _, kernel_fn = stax.serial(
          stax.Dense(1), stax.Elementwise(nngp_fn=nngp_fn)
      )
  
      # Below we construct the kernel using the manually-derived NTK expression.
      _, _, kernel_fn_manual = stax.serial(
          stax.Dense(1), stax.ExpNormalized()
      )
  
      key = random.PRNGKey(1)
      x1 = random.normal(key, (10, 2))
      x2 = random.normal(key, (20, 2))
  
      k_auto = kernel_fn(x1, x2, "ntk")
      k_manual = kernel_fn_manual(x1, x2, "ntk")
  
      # The two kernels match!
      assert jnp.max(jnp.abs(k_manual - k_auto)) < 1e-6
      print("NTK derived via autodiff matches the hand-derived NTK!")
  
imports: from absl import app
  from jax import numpy
  from jax import random
  from neural_tangents import stax
label_desc: Use image node_id 1 for calling a global function with example usage: main(unused_argv) and returns: unknown