code: def test_fast_eval_shape_within_transform(self):
      def f(x):
          m = basic.Linear(20)
          y_slow = stateful.eval_shape(m, x)
          y_fast = eval_shape.fast_eval_shape(m, x)
          self.assertEqual(y_slow, y_fast)
          return m(x)
  
      f = transform.transform(f)
      rng = jax.random.PRNGKey(0)
      x = jnp.ones([1, 12])
      params = jax.eval_shape(f.init, rng, x)
      self.assertEqual(
          params,
          {
              "linear": {
                  "w": jax.ShapeDtypeStruct((12, 20), jnp.float32),
                  "b": jax.ShapeDtypeStruct((20,), jnp.float32),
              }
          },
      )
      jax.eval_shape(f.apply, params, rng, x)
  
file_path: dm-haiku/haiku/_src/eval_shape_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import basic
  from haiku._src import eval_shape
  from haiku._src import stateful
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
label_desc: Use image node_id 4 for calling the EvalShapeTest obj's underlying member method code with example usage: obj.test_fast_eval_shape_within_transform() and returns: m