code: def test_max_pool_batch_vs_vmap(self):
      key = jax.random.PRNGKey(42)
      batch = jax.random.normal(key, [8, 28, 28, 3])
      p = functools.partial(
          pool.max_pool,
          window_shape=(4, 4, 1),
          strides=(2, 2, 1),
          padding="VALID",
      )
      np.testing.assert_allclose(p(batch), jax.vmap(p)(batch))
  
imports: import functools
  from absl.testing import absltest
  from haiku._src import pool
  from haiku._src import test_utils
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling the MaxPoolTest obj's underlying member method code with example usage: obj.test_max_pool_batch_vs_vmap() without return types