code: def test_jitted_f(self):
      witness = []
  
      def f(x):
          witness.append(None)
          return basic.Linear(1)(x)
  
      f = transform.transform(f)
      rng = jax.random.PRNGKey(42)
      x = jnp.zeros([1, 1])
      params = f.init(rng, x)
      del witness[:]
  
      # This layer of indirection (`g`) means summarise cannot unpack `f` and
      # strip our jit.
      jit_apply = jax.jit(f.apply)
      g = lambda params, x: jit_apply(params, None, x)
      for _ in range(2):
          g(params, x)  # Warm up JIT.
          self.assertLen(witness, 1)
  
      summary = get_summary(g, params, x)
      self.assertLen(summary, 1)
  
imports: from collections.abc import Sequence
  import typing
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import base
  from haiku._src import basic
  from haiku._src import module
  from haiku._src import summarise
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
label_desc: Use image node_id 5 for calling the SummariseTest obj's underlying member method code with example usage: obj.test_jitted_f() and returns: basic