code: def test_conv_local_deep(
      self, get, pool, same_inputs, readout, parameterization
  ):
      test_utils.skip_test(self)
  
      key1, key2, key_mc = random.split(random.PRNGKey(1), 3)
      x1 = random.normal(key1, (2, 7, 8, 3))
      x2 = (
          random.normal(key2, (3, 7, 8, 3)) if not same_inputs else None
      )
  
      def get_nn(conv):
          width = 256
          return stax.serial(
              conv(
                  width,
                  (2, 3),
                  (2, 1),
                  padding="CIRCULAR",
                  W_std=1.5,
                  b_std=0.2,
                  parameterization=parameterization,
              ),
              pool,
              stax.Erf(),
              conv(width, (3, 1), (1, 2), padding="SAME"),
              stax.Relu(),
              conv(
                  width,
                  (2, 3),
                  (2, 1),
                  padding="VALID",
                  W_std=1.2,
                  b_std=0.3,
                  parameterization=parameterization,
              ),
              readout,
              stax.Dense(1 if get == "ntk" else width),
          )
  
      init_fn, apply_fn, kernel_fn_local = get_nn(stax.ConvLocal)
  
      k_local = kernel_fn_local(x1, x2, get)
  
      # Test results for consistency with different diagonalizations.
      for diagonal_batch in [True]:
          for diagonal_spatial in [True, False]:
              kwargs = dict(
                  get=get,
                  diagonal_batch=diagonal_batch,
                  diagonal_spatial=diagonal_spatial,
              )
              with self.subTest(**kwargs):
                  k_local_d = kernel_fn_local(x1, x2, **kwargs)
                  test_utils.assert_close_matrices(
                      self, k_local, k_local_d, 0.01
                  )
  
      # Test against CNN-GP diagonal if only flattening is used.
      if (
          pool[0].__name__ == "Identity"
          and readout[0].__name__ == "Flatten"
      ):
          _, _, kernel_fn_conv = get_nn(stax.Conv)
          k_conv = kernel_fn_conv(x1, x2, get)
          self.assertAllClose(k_conv, k_local)
  
      # Test against MC.
      kernel_fn_mc = nt.monte_carlo_kernel_fn(
          init_fn,
          apply_fn,
          key_mc,
          n_samples=512,
          device_count=0,
          implementation=_DEFAULT_TESTING_NTK_IMPLEMENTATION,
          vmap_axes=0,
      )
      k_mc = kernel_fn_mc(x1, x2, get)
      test_utils.assert_close_matrices(self, k_mc, k_local, 0.015, 1.0)
  
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 3 for calling the ConvLocalTest obj's underlying member method code with example usage: obj.test_conv_local_deep(get, pool, same_inputs, readout, parameterization) and returns: stax