code: def test_vmap(self, module_fn: ModuleFn, shape, dtype):
      rng = jax.random.PRNGKey(42)
      if jnp.issubdtype(dtype, jnp.integer):
          x = jax.random.randint(rng, shape, 0, np.prod(shape), dtype)
      else:
          x = jax.random.uniform(rng, shape, dtype)
  
      # Expand our input since we will map over it.
      x = jnp.broadcast_to(x, (2,) + x.shape)
  
      f = hk.transform_with_state(
          lambda x: module_fn()(x)
      )  # pylint: disable=unnecessary-lambda
      f_mapped = hk.transform_with_state(
          lambda x: hk.vmap(lambda x: module_fn()(x), split_rng=False)(
              x
          )
      )  # pylint: disable=unnecessary-lambda
  
      params, state = f_mapped.init(rng, x)
  
      # JAX vmap with explicitly unmapped params/state/rng. This should be
      # equivalent to `f_mapped.apply(..)` (since by default hk.vmap does not map
      # params/state/rng).
      v_apply = jax.vmap(
          f.apply, in_axes=(None, None, None, 0), out_axes=(0, None)
      )
  
      module_type = descriptors.module_type(module_fn)
      atol = CUSTOM_ATOL.get(module_type, DEFAULT_ATOL)
      assert_allclose = functools.partial(
          np.testing.assert_allclose, atol=atol
      )
      jax.tree_util.tree_map(
          assert_allclose,
          f_mapped.apply(params, state, rng, x),
          v_apply(params, state, rng, x),
      )
  
file_path: dm-haiku/haiku/_src/integration/hk_transforms_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import haiku
  from haiku._src import stateful
  from haiku._src import test_utils
  from haiku._src.integration import descriptors
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 5 for calling the HaikuTransformsTest obj's underlying member method code with example usage: obj.test_vmap(module_fn, shape, dtype) without return types