code: def test_conv_local_conv(self):
      test_utils.skip_test(self, platforms=("cpu", "tpu"))
  
      key1, key2 = random.split(random.PRNGKey(1), 2)
      x1 = jnp.cos(random.normal(key1, (5, 32, 32, 1)))
      x2 = jnp.sin(random.normal(key2, (5, 32, 32, 1)))
  
      width = 128
      local_conv = stax.serial(
          stax.ConvLocal(width, (3, 2)),
          stax.AvgPool((2, 3), padding="SAME"),
          stax.Relu(),
          stax.ConvLocal(width, (1, 2), padding="SAME"),
          stax.AvgPool((2, 1), padding="SAME"),
          stax.Relu(),
          stax.Conv(width, (3, 3), padding="SAME"),
          stax.Relu(),
          stax.Conv(width, (3, 3), padding="SAME"),
      )
  
      init_fn, apply_fn, kernel_fn = local_conv
  
      # No projection layer
      k = kernel_fn(x1, x2)
      self.assertEqual(k.diagonal_spatial, False)
      self._test_against_mc(apply_fn, init_fn, k.nngp, x1, x2, 0.03)
  
      # Top layer flat
      init_fn, apply_fn, kernel_fn = stax.serial(
          local_conv, stax.Flatten()
      )
      k_jit = jit(lambda x1, x2: kernel_fn(x1, x2))
      k_jit(x2, x1).nngp.block_until_ready()
      time_flat = time.time()
      k = k_jit(x1, x2).nngp.block_until_ready()
      time_flat = time.time() - time_flat
      self._test_against_mc(apply_fn, init_fn, k, x1, x2, 0.03)
  
      # Top layer pooling
      init_fn, apply_fn, kernel_fn = stax.serial(
          local_conv, stax.GlobalAvgPool()
      )
      k_jit = jit(lambda x1, x2: kernel_fn(x1, x2))
      k_jit(x2, x1).nngp.block_until_ready()
      time_pool = time.time()
      k = k_jit(x1, x2).nngp.block_until_ready()
      time_pool = time.time() - time_pool
      self.assertLess(time_flat * 4, time_pool)
      self._test_against_mc(apply_fn, init_fn, k, x1, x2, 0.03)
  
      # Top layer LCN + pooling
      init_fn, apply_fn, kernel_fn = stax.serial(
          local_conv,
          stax.ConvLocal(width, (2, 2), padding="SAME"),
          stax.GlobalAvgPool(),
      )
      k_jit = jit(lambda x1, x2: kernel_fn(x1, x2))
      k_jit(x2, x1).nngp.block_until_ready()
      time_lcn_pool = time.time()
      k = k_jit(x1, x2).nngp.block_until_ready()
      time_lcn_pool = time.time() - time_lcn_pool
      self.assertLess(time_lcn_pool * 5, time_pool)
      self._test_against_mc(apply_fn, init_fn, k, x1, x2, 0.03)
  
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 4 for calling the ConvLocalTest obj's underlying member method code with example usage: obj.test_conv_local_conv() without return types