code: def test_fast_eval_shape_inside_transform(
      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)
  
      m = module_fn()
      m_slow = hk.eval_shape(m, x)
      m_fast = hk.experimental.fast_eval_shape(m, x)
      self.assertEqual(m_slow, m_fast)
  
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 7 for calling the HaikuTransformsTest obj's underlying member method code with example usage: obj.test_fast_eval_shape_inside_transform(module_fn, shape, dtype) without return types