code: def _test_analytic_kernel_composition(self, batching_fn):
      # Check Fully-Connected.
      rng = random.PRNGKey(0)
      rng_self, rng_other = random.split(rng)
      x_self = random.normal(rng_self, (8, 2))
      x_other = random.normal(rng_other, (2, 2))
      Block = stax.serial(stax.Dense(256), stax.Relu())
  
      _, _, ker_fn = Block
      ker_fn = batching_fn(ker_fn)
  
      _, _, composed_ker_fn = stax.serial(Block, Block)
  
      ker_out = ker_fn(ker_fn(x_self))
      composed_ker_out = composed_ker_fn(x_self)
      if batching_fn == batching._parallel:
          # In the parallel setting, `x1_is_x2` is not computed correctly
          # when x1==x2.
          composed_ker_out = composed_ker_out.replace(
              x1_is_x2=ker_out.x1_is_x2
          )
      self.assertAllClose(ker_out, composed_ker_out)
  
      ker_out = ker_fn(ker_fn(x_self, x_other))
      composed_ker_out = composed_ker_fn(x_self, x_other)
      if batching_fn == batching._parallel:
          composed_ker_out = composed_ker_out.replace(
              x1_is_x2=ker_out.x1_is_x2
          )
      self.assertAllClose(ker_out, composed_ker_out)
  
      # Check convolutional + pooling.
      x_self = random.normal(rng, (8, 4, 4, 3))
      x_other = random.normal(rng, (2, 4, 4, 3))
  
      Block = stax.serial(stax.Conv(256, (2, 2)), stax.Relu())
      Readout = stax.serial(stax.GlobalAvgPool(), stax.Dense(10))
  
      block_ker_fn, readout_ker_fn = Block[2], Readout[2]
      _, _, composed_ker_fn = stax.serial(Block, Readout)
      block_ker_fn = batching_fn(block_ker_fn)
      readout_ker_fn = batching_fn(readout_ker_fn)
  
      ker_out = readout_ker_fn(block_ker_fn(x_self))
      composed_ker_out = composed_ker_fn(x_self)
      if batching_fn == batching._parallel:
          composed_ker_out = composed_ker_out.replace(
              x1_is_x2=ker_out.x1_is_x2
          )
      self.assertAllClose(ker_out, composed_ker_out)
      ker_out = readout_ker_fn(block_ker_fn(x_self, x_other))
      composed_ker_out = composed_ker_fn(x_self, x_other)
      if batching_fn == batching._parallel:
          composed_ker_out = composed_ker_out.replace(
              x1_is_x2=ker_out.x1_is_x2
          )
      self.assertAllClose(ker_out, composed_ker_out)
  
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 6 for calling the BatchTest obj's underlying member method code with example usage: obj._test_analytic_kernel_composition(batching_fn) without return types