docstr: This test creates a Table using `module_summary` and checks that it
  matches the expected output given the CNN model defined in
  `_get_tabulate_cnn`.
docstr_tok: ["This","test","creates","a","Table","using","`","module_summary","`","and","checks","that","it","matches","the","expected","output","given","the","CNN","model","defined","in","`","_get_tabulate_cnn","`","."]
code: def test_module_summary(self):
      """
      This test creates a Table using `module_summary` 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)
  
      table = summary._get_module_table(
          module,
          depth=None,
          show_repeated=True,
          compute_flops=True,
          compute_vjp_flops=True,
      )(
          {"dropout": random.key(0), "params": random.key(1)},
          x,
          training=True,
          mutable=True,
      )
      # get values for inputs and outputs from their _ValueRepresentation
      for row in table:
          row.inputs = jax.tree_util.tree_map(
              _get_obj_repr_value, row.inputs
          )
          row.outputs = jax.tree_util.tree_map(
              _get_obj_repr_value, row.outputs
          )
  
      # 10 rows = 1 CNN + 4 ConvBlock_0 + 4 ConvBlock_1 + 1 Dense_0
      self.assertLen(table, 10)
  
      # check paths
      self.assertEqual(table[0].path, ())
  
      self.assertEqual(table[1].path, ("block1",))
      self.assertEqual(table[2].path, ("block1", "conv"))
      self.assertEqual(table[3].path, ("block1", "bn"))
      self.assertEqual(table[4].path, ("block1", "dropout"))
  
      self.assertEqual(table[5].path, ("block2",))
      self.assertEqual(table[6].path, ("block2", "conv"))
      self.assertEqual(table[7].path, ("block2", "bn"))
      self.assertEqual(table[8].path, ("block2", "dropout"))
  
      self.assertEqual(table[9].path, ("dense",))
  
      # check outputs shapes
      self.assertEqual(
          (table[0].inputs[0].shape, table[0].inputs[1]),
          (x.shape, dict(training=True)),
      )
      self.assertEqual(
          _get_shapes(table[0].outputs),
          (
              (batch_size, 10),
              dict(a=(batch_size, 10), b=(batch_size, 10)),
          ),
      )
  
      self.assertEqual(
          _get_shapes(table[1].inputs),
          ((batch_size, 28, 28, 1), {"training": True}),
      )
      self.assertEqual(table[1].outputs.shape, (batch_size, 28, 28, 32))
      self.assertEqual(table[2].inputs.shape, (batch_size, 28, 28, 1))
      self.assertEqual(table[2].outputs.shape, (batch_size, 28, 28, 32))
      self.assertEqual(
          _get_shapes(table[3].inputs),
          ((batch_size, 28, 28, 32), {"use_running_average": False}),
      )
      self.assertEqual(table[3].outputs.shape, (batch_size, 28, 28, 32))
      self.assertEqual(
          _get_shapes(table[4].inputs),
          ((batch_size, 28, 28, 32), {"deterministic": False}),
      )
      self.assertEqual(table[4].outputs.shape, (batch_size, 28, 28, 32))
  
      self.assertEqual(
          _get_shapes(table[5].inputs),
          ((batch_size, 28, 28, 32), {"training": True}),
      )
      self.assertEqual(table[5].outputs.shape, (batch_size, 28, 28, 64))
      self.assertEqual(table[6].inputs.shape, (batch_size, 28, 28, 32))
      self.assertEqual(table[6].outputs.shape, (batch_size, 28, 28, 64))
      self.assertEqual(
          _get_shapes(table[7].inputs),
          ((batch_size, 28, 28, 64), {"use_running_average": False}),
      )
      self.assertEqual(table[7].outputs.shape, (batch_size, 28, 28, 64))
      self.assertEqual(
          _get_shapes(table[8].inputs),
          ((batch_size, 28, 28, 64), {"deterministic": False}),
      )
      self.assertEqual(table[8].outputs.shape, (batch_size, 28, 28, 64))
  
      self.assertEqual(table[9].inputs.shape, (batch_size, 64))
      self.assertEqual(table[9].outputs.shape, (batch_size, 10))
  
      # check no summary is performed
      for row in table:
          self.assertEqual(
              row.module_variables,
              row.counted_variables,
          )
  
      # Each module FLOPs >= sum of its submodule FLOPs.
      # Can be greater due to ops like `nn.relu` not belonging to any submodule.
      for r in table:
          flops, vjp_flops = r.flops, r.vjp_flops
          submodule_flops, submodule_vjp_flops = 0, 0
          for s in table:
              if (
                  len(s.path) == len(r.path) + 1
                  and s.path[: len(r.path)] == r.path
              ):
                  submodule_flops += s.flops
                  submodule_vjp_flops += s.vjp_flops
  
          self.assertGreaterEqual(flops, submodule_flops)
          self.assertGreaterEqual(vjp_flops, submodule_vjp_flops)
  
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 1 for calling the SummaryTest obj's underlying member method code with example usage: obj.test_module_summary() without return types