docstr: This test creates a string representation of a Module using
  `Module.tabulate` and checks that it matches the expected output given the
  CNN model defined in `_get_tabulate_cnn`.
docstr_tok: ["This","test","creates","a","string","representation","of","a","Module","using","`","Module.tabulate","`","and","checks","that","it","matches","the","expected","output","given","the","CNN","model","defined","in","`","_get_tabulate_cnn","`","."]
code: def test_tabulate_function(self):
      """
      This test creates a string representation of a Module using
      `Module.tabulate` and checks that it matches the expected output given the
      CNN model defined in `_get_tabulate_cnn`.
      """
      batch_size = 32
  
      x = jnp.ones((batch_size, 28, 28, 1))
      module = CNN(test_sow=False)
  
      module_repr = nn.tabulate(
          module,
          {"dropout": random.key(0), "params": random.key(1)},
          console_kwargs=CONSOLE_TEST_KWARGS,
          compute_flops=True,
          compute_vjp_flops=True,
      )(x, training=True)
  
      lines = module_repr.split("\n")
  
      # check title
      module_name = module.__class__.__name__
      self.assertIn(f"{module_name} Summary", lines[1])
  
      # check headers are correct
      self.assertIn("path", lines[3])
      self.assertIn("module", lines[3])
      self.assertIn("inputs", lines[3])
      self.assertIn("outputs", lines[3])
      self.assertIn("params", lines[3])
      self.assertIn("flops", lines[3])
      self.assertIn("batch_stats", lines[3])
  
      # collection counts
      self.assertIn("Total", lines[-6])
      self.assertIn("192", lines[-6])
      self.assertIn("768 B", lines[-6])
      self.assertIn("19,658", lines[-6])
      self.assertIn("78.6 KB", lines[-6])
  
      # total counts
      self.assertIn("Total Parameters", lines[-3])
      self.assertIn("19,850", lines[-3])
      self.assertIn("79.4 KB", lines[-3])
  
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 6 for calling the SummaryTest obj's underlying member method code with example usage: obj.test_tabulate_function() without return types