docstr: Tests forward- and reverse-mode autodiff for nonlinearities.
code: def test_activations(
      self,
      get,
      parameterization,
      parameterization_out,
      x1_type,
      x2_type,
      b_std,
      phi,
      do_jit,
  ):
      """Tests forward- and reverse-mode autodiff for nonlinearities."""
      if phi == stax.ABRelu:
          phi_ = phi(0.25, 0.5)
      else:
          phi_ = phi()
  
      if phi not in [stax.Relu]:
          test_utils.skip_test(self)
  
      n_out = 1 if get == "ntk" else 1024
      width = 612
  
      W_std_in = (
          width ** (-0.5) if parameterization_out == "standard" else 1.0
      )
      if phi == stax.Exp:
          W_std_in /= 10.0
  
      init_fn, apply_fn, kernel_fn = stax.serial(
          stax.Dense(
              width,
              W_std=W_std_in,
              b_std=b_std,
              parameterization=parameterization,
          ),
          phi_,
          stax.Dense(
              n_out, b_std=b_std, parameterization=parameterization_out
          ),
      )
  
      def get_x(x_type, key):
          shape = (1, 2)
          if x_type == "zeros":
              x = jnp.zeros(shape)
          elif x_type == "ones":
              x = jnp.ones(shape)
          elif x_type == "random":
              x = random.normal(random.PRNGKey(key), shape)
          elif x_type == "sin":
              x = jnp.sin(random.normal(random.PRNGKey(key), shape))
          elif x_type == "none":
              return None
          else:
              raise ValueError(x_type)
          return x
  
      x1 = get_x(x1_type, 1)
      if x2_type == "x1":
          x2 = x1
      else:
          x2 = get_x(x2_type, 2)
  
      def kernel_scalar(x1, x2):
          return kernel_fn(x1, x2, get)[0, 0]
  
      if do_jit:
          kernel_scalar = jit(kernel_scalar)
  
      k1 = kernel_scalar(x1, x2)
      k2, k2_grad = value_and_grad(kernel_scalar)(x1, x2)
      self.assertAllClose(k1, k2)
  
      # Compare to forward-mode.
      k2_fwd, _ = jvp(kernel_scalar, (x1, x2), (x1, x2))
      k2_grad_fwd = jacfwd(kernel_scalar)(x1, x2)
      self.assertAllClose(k1, k2_fwd)
      self.assertAllClose(k2_grad, k2_grad_fwd)
  
      # `stax.ExpNormalized` has no forward pass.
      # `stax.Sign` is discontinuous at `0`, so NTK MC kernel does not converge to
      # infinite-width kernel.
      if phi == stax.ExpNormalized or (
          get == "ntk" and phi == stax.Sign
      ):
          raise absltest.SkipTest("Not comparing against MC kernels.")
  
      _kernel_scalar_mc = nt.monte_carlo_kernel_fn(
          init_fn,
          apply_fn,
          key=random.PRNGKey(3),
          n_samples=1,
          device_count=0,
      )
  
      def kernel_scalar_mc(x1, x2):
          return _kernel_scalar_mc(x1, x2, get)[0, 0]
  
      k_mc = kernel_scalar_mc(x1, x2)
      k_mc2, k_mc2_grad = value_and_grad(kernel_scalar_mc)(x1, x2)
      self.assertAllClose(k_mc, k_mc2)
  
      # Compare MC to forward-mode.
      k_mc2_fwd, _ = jvp(kernel_scalar_mc, (x1, x2), (x1, x2))
      k_mc2_grad_fwd = jacfwd(kernel_scalar_mc)(x1, x2)
      self.assertAllClose(k_mc, k_mc2_fwd)
      self.assertAllClose(k_mc2_grad, k_mc2_grad_fwd)
  
      def kernel_fn_emp(x1, x2, get, params):
          return nt.empirical_kernel_fn(apply_fn)(x1, x2, get, params)[
              0, 0
          ]
  
      kernel_fn_emp_g = jit(
          value_and_grad(kernel_fn_emp), static_argnames="get"
      )
  
      def kernel_scalar_mc_grad_mean(x1, x2):
          x1: jnp.ndarray
          key = random.PRNGKey(4)
          n_samples = 2**9
          k, k_grad = 0.0, 0.0
  
          for _ in range(n_samples):
              _, params = init_fn(key, x1.shape)
              k_mc2, k_mc2_grad = kernel_fn_emp_g(x1, x2, get, params)
              k += k_mc2
              k_grad += k_mc2_grad
              key, _ = random.split(key)
  
          k /= n_samples
          k_grad /= n_samples
          return k, k_grad
  
      k_mc2_mean, k_mc2_grad_mean = kernel_scalar_mc_grad_mean(x1, x2)
  
      # Compare kernels.
      self.assertAllClose(k1, k_mc2_mean, atol=8e-3, rtol=4e-2)
  
      if phi == stax.Sign and get == "nngp":
          raise absltest.SkipTest(
              "Derivative of the empirical NNGP of a "
              "discontinuous function does not converge "
              "to the derivative of the infinite width NNGP."
          )
  
      if (
          phi in [stax.Abs, stax.Relu, stax.LeakyRelu, stax.ABRelu]
          and get == "ntk"
      ):
          raise absltest.SkipTest(
              "Derivative of the empirical NTK of a "
              "non-differentiable function does not converge "
              "to the derivative of the infinite width NTK."
          )
  
      atol = 1e-2
  
      # Compare gradient of the analytic kernel to empirical kernel.
      if jnp.max(jnp.abs(k2_grad - k_mc2_grad_mean)) > atol:
          test_utils.assert_close_matrices(
              self, k_mc2_grad_mean, k2_grad, rtol=0.05, atol=10.0
          )
  
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 2 for calling the AutodiffTest obj's underlying member method code with example usage: obj.test_activations(get, parameterization, parameterization_out, x1_type, x2_type, b_std, phi, do_jit) and returns: x, kernel_fn, _kernel_scalar_mc, nt, k,  k_grad, None