code: def _theoretical_kernel(
      unused_key, input_shape, network, just_theta, use_dropout
  ):
      _, _, _kernel_fn = _build_network(
          input_shape, network, 1, use_dropout
      )
  
      @jit
      def kernel_fn(x1, x2=None):
          get_all = None
          k = (
              _kernel_fn(x1, x2, "ntk")
              if just_theta
              else _kernel_fn(x1, x2, get_all)
          )
          return k
  
      return kernel_fn
  
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 3 for calling a global function with example usage: _theoretical_kernel(unused_key, input_shape, network, just_theta, use_dropout) and returns: kernel_fn, k