code: def _test_activation_fc(self, phi, get):
      key1, key2, key_mc = random.split(random.PRNGKey(1), 3)
      x1 = jnp.cos(random.normal(key1, (3, 2)))
      x2 = jnp.cos(random.normal(key2, (2, 2)))
  
      init_fn, apply_fn, kernel_fn = stax.serial(
          stax.Dense(1024), phi, stax.Dense(1 if get == "ntk" else 1024)
      )
  
      analytic_kernel = kernel_fn(x1, x2, get, diagonal_spatial=True)
      mc_kernel_fn = nt.monte_carlo_kernel_fn(
          init_fn=init_fn,
          apply_fn=apply_fn,
          key=key_mc,
          n_samples=800,
          implementation=2,
          vmap_axes=0,
          device_count=0,
      )
  
      if get == "cov1":
          empirical_kernel = jnp.diag(mc_kernel_fn(x1, None, "nngp"))
      else:
          empirical_kernel = mc_kernel_fn(x1, x2, get)
  
      self.assertAllClose(
          analytic_kernel, empirical_kernel, atol=0.01, rtol=0.05
      )
  
imports: import itertools
  import random
  from absl.testing import absltest
  from jax import config
  from jax import default_backend
  from jax import grad, jacfwd, jacrev, jit, jvp, value_and_grad
  from jax import random
  import jax.numpy
  import neural_tangents
  from neural_tangents import stax
  from neural_tangents._src.empirical import _DEFAULT_TESTING_NTK_IMPLEMENTATION
  from tests import test_utils
label_desc: Use image node_id 1 for calling the ActivationTest obj's underlying member method code with example usage: obj._test_activation_fc(phi, get) without return types