code: def _empirical_kernel(
      key, input_shape, network, out_logits, use_dropout
  ):
      init_fn, f, _ = _build_network(
          input_shape, network, out_logits, use_dropout
      )
      key, split = random.split(key)
      _, params = init_fn(key, (-1,) + input_shape)
      kernel_fn = jit(
          nt.empirical_ntk_fn(
              f, implementation=_DEFAULT_TESTING_NTK_IMPLEMENTATION
          )
      )
      return partial(kernel_fn, params=params, keys=split)
  
imports: from functools import partial
  from absl.testing import absltest
  from jax import jit
  from jax import config
  from jax import random
  import jax.numpy
  from jax.tree_util import tree_map
  import neural_tangents
  from neural_tangents import stax
  from neural_tangents._src import batching
  from neural_tangents._src.empirical import _DEFAULT_TESTING_NTK_IMPLEMENTATION
  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, use_dropout) and returns: partial