code: def _test_kernel_against_batched(
      cls,
      kernel_fn,
      batched_kernel_fn,
      train,
      test,
      is_parallel_only=False,
  ):
      g = kernel_fn(train, None)
      g_b = batched_kernel_fn(train, None)
  
      if is_parallel_only and hasattr(g_b, "x1_is_x2"):
          # In the parallel setting, `x1_is_x2` is not computed correctly when x1==x2.
          g_b = g_b.replace(x1_is_x2=g.x1_is_x2)
  
      cls.assertAllClose(g, g_b)
  
      g = kernel_fn(train, test)
      g_b = batched_kernel_fn(train, test)
      if is_parallel_only and hasattr(g_b, "x1_is_x2"):
          g_b = g_b.replace(x1_is_x2=g.x1_is_x2)
      cls.assertAllClose(g, g_b)
  
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 4 for calling a global function with example usage: _test_kernel_against_batched(cls, kernel_fn, batched_kernel_fn, train, test, is_parallel_only) without return types