code: def _get_ntks(f, x1, x2, params, vmap_axes):
      """Returns a list of NTKs computed using different implementations."""
      kwargs = dict(
          f=f,
          trace_axes=(),
          vmap_axes=vmap_axes,
      )
  
      # Default, baseline Jacobian contraction.
      jacobian_contraction = nt.experimental.empirical_ntk_fn_tf(
          **kwargs,
          implementation=nt.NtkImplementation.JACOBIAN_CONTRACTION,
      )
      # (6, 3, 10, 10) full `jnp.ndarray` test-train NTK
      ntk_jc = jacobian_contraction(x2, x1, params)
  
      # NTK-vector products-based implementation.
      ntk_vector_products = nt.experimental.empirical_ntk_fn_tf(
          **kwargs,
          implementation=nt.NtkImplementation.NTK_VECTOR_PRODUCTS,
      )
      ntk_vp = ntk_vector_products(x2, x1, params)
  
      # Structured derivatives-based implementation.
      structured_derivatives = nt.experimental.empirical_ntk_fn_tf(
          **kwargs,
          implementation=nt.NtkImplementation.STRUCTURED_DERIVATIVES,
      )
      ntk_sd = structured_derivatives(x2, x1, params)
  
      # Auto-FLOPs-selecting implementation. Doesn't work correctly on CPU/GPU.
      auto = nt.experimental.empirical_ntk_fn_tf(
          **kwargs, implementation=nt.NtkImplementation.AUTO
      )
      ntk_auto = auto(x2, x1, params)
  
      return [ntk_jc, ntk_vp, ntk_sd, ntk_auto]
  
imports: from absl import app
  import neural_tangents
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: _get_ntks(f, x1, x2, params, vmap_axes) and returns: list