code: def _compare_ntks(
      self, f, f_jax, params, trace_axes, diagonal_axes, vmap_axes
  ):
      if any(i == j for i in trace_axes for j in diagonal_axes):
          raise absltest.SkipTest(
              "Overlapping trace and diagonal axes."
          )
  
      kwargs = dict(
          trace_axes=trace_axes,
          diagonal_axes=diagonal_axes,
      )
  
      jax_ntk_fns = [
          jax.jit(
              nt.empirical_ntk_fn(
                  **kwargs, f=f_jax, implementation=i, vmap_axes=v
              )
          )
          for i in nt.NtkImplementation
          for v in vmap_axes
          if v not in trace_axes + diagonal_axes
      ]
  
      ntk_fns = [
          experimental.empirical_ntk_fn_tf(
              **kwargs, f=f, implementation=i, vmap_axes=v
          )
          for i in nt.NtkImplementation
          for v in vmap_axes
          if v not in trace_axes + diagonal_axes
      ]
  
      x_shape = (
          f.input_shape[1:]
          if isinstance(f, tf.Module)
          else f.input_signature[1].shape[1:]
      )
  
      x1 = (
          tf.random.normal((2,) + x_shape, seed=2)
          / np.prod(x_shape) ** 0.5
      )
      x2 = (
          tf.random.normal((3,) + x_shape, seed=3)
          / np.prod(x_shape) ** 0.5
      )
  
      x1_jax = jnp.array(x1)
      x2_jax = jnp.array(x2)
      params_jax = jax.tree_map(jnp.array, params)
  
      jax_ntks = [
          ntk_fn_i(x1_jax, x2_jax, params_jax)
          for ntk_fn_i in jax_ntk_fns
      ]
  
      ntks = list(
          enumerate([ntk_fn_i(x1, x2, params) for ntk_fn_i in ntk_fns])
      )
  
      if len(tf.config.list_physical_devices()) > 1:  # TPU
          atol = 0.0
          rtol = 5e-3
          atol_jax = 0.4
          rtol_jax = 0.15  # TODO(romann): revisit poor TPU agreement.
      else:
          atol = 1e-5
          rtol = 1e-4
          atol_jax = 0.0
          rtol_jax = 5e-5
  
      for i1, ntk1 in ntks:
          for i2, ntk2 in ntks[i1 + 1 :]:
              # Compare different implementation
              np.testing.assert_allclose(
                  ntk1, ntk2, rtol=rtol, atol=atol
              )
              # Compare against the JAX version (without calling `jax2tf`).
              np.testing.assert_allclose(
                  ntk1, jax_ntks[i1], rtol=rtol_jax, atol=atol_jax
              )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import numpy
  import neural_tangents
  from neural_tangents import experimental
  import numpy
  import tensorflow
label_desc: Use image node_id 1 for calling the EmpiricalTfTest obj's underlying member method code with example usage: obj._compare_ntks(f, f_jax, params, trace_axes, diagonal_axes, vmap_axes) without return types