code: def test_compute_flops(self):
      def _compute_flops(
          eqn: jax.core.JaxprEqn, expression: jaxpr_info.Expression
      ) -> int:
          del expression
          return max(
              np.prod(var.aval.shape) for var in eqn.invars
          )  # pytype: disable=attribute-error
  
      def add(x, y):
          return jnp.sign(x) + jnp.cos(y)
  
      a = jnp.zeros((12, 7))
      mod = jaxpr_info.make_model_info(
          add, compute_flops=_compute_flops
      )(a, a)
      self.assertContentsEqual(
          jaxpr_info.format_module(mod),
          """
  add 252 flops
    sign 84 flops
      sign 84 flops in f32[12,7], out f32[12,7]
    cos 84 flops in f32[12,7], out f32[12,7]
    add 84 flops in f32[12,7], f32[12,7], out f32[12,7]
  """,
      )
  
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 2 for calling the JaxprInfoTest obj's underlying member method code with example usage: obj.test_compute_flops() and returns: max, unknown