code: def test_vmap(self, module_fn: ModuleFn, shape, dtype):
      batch_size, shape = shape[0], shape[1:]
      rng = jax.random.PRNGKey(42)
      if jnp.issubdtype(dtype, jnp.integer):
          sample = jax.random.randint(
              rng, shape, 0, np.prod(shape), dtype
          )
      else:
          sample = jax.random.uniform(rng, shape, dtype)
      batch = jnp.broadcast_to(sample, (batch_size,) + sample.shape)
  
      def g(x):
          return module_fn()(x)
  
      f = hk.transform_with_state(g)
  
      # Ensure application under vmap is the same.
      params, state = f.init(rng, sample)
      v_apply = jax.vmap(f.apply, in_axes=(None, None, None, 0))
      jax.tree_util.tree_map(
          lambda a, b: np.testing.assert_allclose(
              a, b, atol=DEFAULT_ATOL
          ),
          f.apply(params, state, rng, batch),
          v_apply(params, state, rng, batch),
      )
  
file_path: dm-haiku/haiku/_src/integration/jax_transforms_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import haiku
  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 2 for calling the JaxTransformsTest obj's underlying member method code with example usage: obj.test_vmap(module_fn, shape, dtype) and returns: module_fn