code: def _check_agreement_with_empirical(
      self, net, same_inputs, use_dropout, is_ntk, rtol=RTOL, atol=ATOL
  ):
      (
          (init_fn, apply_fn, kernel_fn),
          input_shape,
          device_count,
          channel_axis,
      ) = net
  
      num_samples = N_SAMPLES * 5 if use_dropout else N_SAMPLES
      key = random.PRNGKey(1)
      x1, x2 = _get_inputs(key, same_inputs, input_shape)
      if default_backend() == "tpu" and use_dropout:
          # including a test case for tpu + dropout with (parallel + batching)
          batch_size = 2
      else:
          batch_size = 0
      x1_out_shape, params = init_fn(key, x1.shape)
      if same_inputs:
          assert x2 is None
      if x2 is None:
          x2_out_shape = x1_out_shape
      else:
          x2_out_shape, params = init_fn(key, x2.shape)
      del params
  
      def _get_empirical(n_samples, get):
          kernel_fn_empirical = nt.monte_carlo_kernel_fn(
              init_fn=init_fn,
              apply_fn=apply_fn,
              key=key,
              n_samples=n_samples,
              device_count=device_count,
              trace_axes=(channel_axis,),
              batch_size=batch_size,
              implementation=_DEFAULT_TESTING_NTK_IMPLEMENTATION,
          )
          if same_inputs:
              assert x2 is None
          return kernel_fn_empirical(x1, x2, get)
  
      if is_ntk:
          exact, shape1, shape2 = kernel_fn(
              x1, x2, ("ntk", "shape1", "shape2")
          )
          empirical = _get_empirical(num_samples, "ntk")
      else:
          exact, shape1, shape2 = kernel_fn(
              x1, x2, ("nngp", "shape1", "shape2")
          )
          empirical = _get_empirical(num_samples, "nngp")
      test_utils.assert_close_matrices(
          self, exact, empirical, rtol, atol
      )
      self.assertEqual(shape1, x1_out_shape)
      self.assertEqual(shape2, x2_out_shape)
  
imports: import functools
  import itertools
  import random
  from absl.testing import absltest
  from jax import config
  from jax import default_backend
  from jax import jit
  from jax import random
  from jax.example_libraries import stax
  import jax.numpy
  import neural_tangents
  from neural_tangents import stax
  from neural_tangents._src.empirical import _DEFAULT_TESTING_NTK_IMPLEMENTATION
  from tests import test_utils
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: _check_agreement_with_empirical(net, same_inputs, use_dropout, is_ntk, rtol, atol) and returns: kernel_fn_empirical