code: def test_tabulate_param_count_and_flops(self):
      class Foo(nn.Module):
          @nn.compact
          def __call__(self, x):
              h = nn.Dense(4)(x)
              return nn.Dense(2)(h)
  
      module = Foo()
      rng = jax.random.key(0)
      x = jnp.ones((16, 9))
  
      rep = module.tabulate(
          rng,
          x,
          console_kwargs=CONSOLE_TEST_KWARGS,
          compute_flops=True,
          compute_vjp_flops=True,
      )
      lines = rep.splitlines()
      self.assertIn("Total Parameters: 50", lines[-2])
  
imports: import enum
  from typing import List
  import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest
  from jax import random
  from flax import linen
  from flax import struct
  from flax.core.scope import Array
  from flax.linen import summary
label_desc: Use image node_id 14 for calling the SummaryTest obj's underlying member method code with example usage: obj.test_tabulate_param_count_and_flops() and returns: nn