code: def test_diagonal_spatial(self, same_inputs, readin, readout):
      kernel_fn, x1, x2 = self._get_kernel_fn(
          same_inputs, readin, readout
      )
      K = kernel_fn(x1, x2)
      K_full = kernel_fn(x1, x2, diagonal_spatial=False)
      batch_shape = x1.shape[0], (x1 if x2 is None else x2).shape[0]
      names = readout[0].__name__, readin[0].__name__
  
      if "GlobalAvgPool" in names:
          if (
              readout[0].__name__ == "GlobalAvgPool"
              and readin[0].__name__ == "Identity"
          ):
              self.assertRaises(
                  ValueError, kernel_fn, x1, x2, diagonal_spatial=True
              )
          self.assertEqual(K_full.nngp.shape, batch_shape)
          self.assertAllClose(K_full, K)
  
      else:
          K_diag = kernel_fn(x1, x2, diagonal_spatial=True)
          if "Flatten" in names:
              self.assertEqual(K_diag.nngp.shape, batch_shape)
              self.assertAllClose(K_diag, K)
              self.assertAllClose(K_diag, K_full)
  
          else:
              self.assertEqual(
                  K_diag.nngp.shape, batch_shape + x1.shape[1:-1]
              )
              self.assertAllClose(K_full, K)
              self.assertAllClose(
                  K_diag.nngp, jnp.einsum("...iijj->...ij", K.nngp)
              )
  
imports: 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
  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
label_desc: Use image node_id 3 for calling the DiagonalTest obj's underlying member method code with example usage: obj.test_diagonal_spatial(same_inputs, readin, readout) without return types