code: def test_avg_pool(self):
      X1 = jnp.ones((4, 2, 3, 2))
      X2 = jnp.ones((3, 2, 3, 2))
  
      _, apply_fn, kernel_fn = stax.AvgPool(
          (2, 2), (1, 1), "SAME", normalize_edges=False
      )
      _, apply_fn_norm, kernel_fn_norm = stax.AvgPool(
          (2, 2), (1, 1), "SAME", normalize_edges=True
      )
      _, apply_fn_stax = ostax.AvgPool((2, 2), (1, 1), "SAME")
  
      out1 = apply_fn((), X1)
      out2 = apply_fn((), X2)
  
      out1_norm = apply_fn_norm((), X1)
      out2_norm = apply_fn_norm((), X2)
  
      out1_stax = apply_fn_stax((), X1)
      out2_stax = apply_fn_stax((), X2)
  
      self.assertAllClose(
          (out1_stax, out2_stax), (out1_norm, out2_norm)
      )
  
      out_unnorm = jnp.array(
          [[1.0, 1.0, 0.5], [0.5, 0.5, 0.25]]
      ).reshape((1, 2, 3, 1))
      out1_unnormalized = jnp.broadcast_to(out_unnorm, X1.shape)
      out2_unnormalized = jnp.broadcast_to(out_unnorm, X2.shape)
  
      self.assertAllClose(
          (out1_unnormalized, out2_unnormalized), (out1, out2)
      )
  
      ker = kernel_fn(X1, X2)
      ker_norm = kernel_fn_norm(X1, X2)
  
      self.assertAllClose(jnp.ones_like(ker_norm.nngp), ker_norm.nngp)
      self.assertAllClose(jnp.ones_like(ker_norm.cov1), ker_norm.cov1)
      self.assertAllClose(jnp.ones_like(ker_norm.cov2), ker_norm.cov2)
  
      self.assertEqual(ker_norm.nngp.shape, ker.nngp.shape)
      self.assertEqual(ker_norm.cov1.shape, ker.cov1.shape)
      self.assertEqual(ker_norm.cov2.shape, ker.cov2.shape)
  
      ker_unnorm = jnp.outer(out_unnorm, out_unnorm).reshape(
          (2, 3, 2, 3)
      )
      ker_unnorm = jnp.transpose(ker_unnorm, axes=(0, 2, 1, 3))
      nngp = jnp.broadcast_to(
          ker_unnorm.reshape((1, 1) + ker_unnorm.shape), ker.nngp.shape
      )
      cov1 = jnp.broadcast_to(
          jnp.expand_dims(ker_unnorm, 0), ker.cov1.shape
      )
      cov2 = jnp.broadcast_to(
          jnp.expand_dims(ker_unnorm, 0), ker.cov2.shape
      )
      self.assertAllClose(
          (nngp, cov1, cov2), (ker.nngp, ker.cov1, ker.cov2)
      )
  
imports: import functools
  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
  from jax.example_libraries import stax
  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
  import numpy
label_desc: Use image node_id 5 for calling the StaxTest obj's underlying member method code with example usage: obj.test_avg_pool() without return types