code: def test_haiku_module(self):
      def forward(x):
          return MyModel()(x)
  
      forward_t = transform.transform_with_state(forward)
  
      rng = jax.random.PRNGKey(42)
      x = jnp.zeros((16, 8, 8, 32))
      params, state = forward_t.init(rng, x)
  
      mod = jaxpr_info.make_model_info(forward_t.apply)(
          params, state, rng, x
      )
      self.assertContentsEqual(
          jaxpr_info.format_module(mod),
          """
  apply_fn
    my_model 4.624 kparams
      conv2_d 4.624 kparams
        conv_general_dilated in f32[16,8,8,32], f32[3,3,32,16], out f32[16,8,8,16]
        broadcast_in_dim in f32[16], out f32[16,8,8,16]
        add in f32[16,8,8,16], f32[16,8,8,16], out f32[16,8,8,16]
  """,
      )
  
imports: from typing import Optional
  from absl import logging
  from absl.testing import absltest
  from haiku._src import conv
  from haiku._src import jaxpr_info
  from haiku._src import module
  from haiku._src import transform
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling the JaxprInfoTest obj's underlying member method code with example usage: obj.test_haiku_module() and returns: MyModel