code: def test_fast_eval_shape(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)
  
      def g(x):
          return module_fn()(x)
  
      f = hk.transform_with_state(g)
  
      init_jax = jax.eval_shape(f.init, rng, x)
      init_hk = hk.experimental.fast_eval_shape(f.init, rng, x)
      self.assertEqual(init_jax, init_hk)
  
      apply_jax = jax.eval_shape(f.apply, *init_jax, rng, x)
      apply_hk = hk.experimental.fast_eval_shape(
          f.apply, *init_hk, rng, x
      )
      self.assertEqual(apply_jax, apply_hk)
  
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 6 for calling the HaikuTransformsTest obj's underlying member method code with example usage: obj.test_fast_eval_shape(module_fn, shape, dtype) and returns: module_fn