code: def _empirical_kernel(key, input_shape, network, out_logits):
      init_fn, f, _ = _build_network(input_shape, network, out_logits)
      _, params = init_fn(key, (-1,) + input_shape)
      _kernel_fn = nt.empirical_kernel_fn(f, trace_axes=(), vmap_axes=0)
      kernel_fn = lambda x1, x2, get: _kernel_fn(x1, x2, get, params)
      return params, f, jit(kernel_fn, static_argnames="get")
  
imports: import math
  from absl.testing import absltest
  from jax import grad
  from jax import jit
  from jax import random
  from jax import vmap
  from jax import config
  from jax.example_libraries import optimizers
  from jax.flatten_util import ravel_pytree
  import jax.numpy
  import jax.tree_util
  import neural_tangents
  from neural_tangents import predict, stax
  from neural_tangents._src.predict import _is_on_cpu
  from tests import test_utils
label_desc: Use image node_id 2 for calling a global function with example usage: _empirical_kernel(key, input_shape, network, out_logits) and returns: params,  f,  jit