code: def test_conv_local(
      self,
      same_inputs,
      padding,
      filter_shape,
      strides,
      size,
      diagonal_batch,
      diagonal_spatial,
      get,
      parameterization,
  ):
      test_utils.skip_test(self)
  
      if diagonal_batch and get != "cov1":
          raise absltest.SkipTest(
              "Checking `diagonal_batch` only on `cov1`."
          )
  
      key1, key2, key_mc = random.split(random.PRNGKey(1), 3)
      shape = (size, 1)
      x1 = random.normal(key1, (2,) + shape)
      x2 = (
          random.normal(key2, (3,) + shape) if not same_inputs else None
      )
  
      kernel_kwargs = dict(
          diagonal_batch=diagonal_batch,
          diagonal_spatial=diagonal_spatial,
      )
  
      conv_kwargs = dict(
          out_chan=512,
          filter_shape=(filter_shape,),
          strides=(strides,),
          padding=padding,
          b_std=0.2,
          W_std=1.5,
          parameterization=parameterization,
      )
  
      init_fn, apply_fn, kernel_fn = stax.ConvLocal(**conv_kwargs)
      k = kernel_fn(x1, x2, **kernel_kwargs)
  
      # Compared to MC estimate
      diagonal_axes = ()
      if diagonal_batch:
          diagonal_axes += (0,)
      if diagonal_spatial:
          diagonal_axes += (1,)
  
      kernel_fn_mc = nt.monte_carlo_kernel_fn(
          init_fn,
          apply_fn,
          key_mc,
          n_samples=512,
          diagonal_axes=diagonal_axes,
          device_count=0,
          implementation=_DEFAULT_TESTING_NTK_IMPLEMENTATION,
          vmap_axes=0,
      )
      k_mc = kernel_fn_mc(
          x1,
          None if get == "cov1" else x2,
          "nngp" if get == "cov1" else get,
      )
      test_utils.assert_close_matrices(
          self, k_mc, getattr(k, get), 0.011, 1.0
      )
  
      # Compared diagonal entries to CNN
      _, _, kernel_fn_conv = stax.Conv(**conv_kwargs)
      k_conv = kernel_fn_conv(x1, x2, **kernel_kwargs)
  
      if not diagonal_spatial:
  
          def get_diag(k):
              k = getattr(k, get)
              k = jnp.diagonal(k, axis1=-1, axis2=-2)
              return k
  
          k_conv = get_diag(k_conv)
          k = get_diag(k)
  
      tol = 0.005 if default_backend() == "tpu" else 0.001
      self.assertAllClose(k_conv, k, atol=tol, rtol=tol)
  
imports: import itertools
  import random
  import string
  import time
  from absl.testing import absltest
  from jax import default_backend
  from jax import jit
  from jax import lax
  from jax import random
  from jax import vjp
  from jax import config
  import jax.numpy
  import more_itertools
  import neural_tangents
  from neural_tangents import stax
  from neural_tangents._src.empirical import _DEFAULT_TESTING_NTK_IMPLEMENTATION
  from neural_tangents._src.utils import utils
  from tests import test_utils
  import numpy
label_desc: Use image node_id 2 for calling the ConvLocalTest obj's underlying member method code with example usage: obj.test_conv_local(same_inputs, padding, filter_shape, strides, size, diagonal_batch, diagonal_spatial, get, parameterization) and returns: k