docstr: This test creates a Table using `module_summary` set the `depth` argument
  to `1`, table should have fewer rows as a consequence.
docstr_tok: ["This","test","creates","a","Table","using","`","module_summary","`","set","the","`","depth","`","argument","to","`","1","`",",","table","should","have","fewer","rows","as","a","consequence","."]
code: def test_module_summary_with_depth(self):
      """
      This test creates a Table using `module_summary` set the `depth` argument
      to `1`, table should have fewer rows as a consequence.
      """
      batch_size = 32
  
      x = jnp.ones((batch_size, 28, 28, 1))
      module = CNN(test_sow=False)
  
      table = summary._get_module_table(
          module,
          depth=1,
          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
          )
  
      # 4 rows = 1 CNN + 1 ConvBlock_0 + 1 ConvBlock_1 + 1 Dense_0
      self.assertLen(table, 4)
  
      # check paths
      self.assertEqual(table[0].path, ())
  
      self.assertEqual(table[1].path, ("block1",))
      self.assertEqual(table[2].path, ("block2",))
      self.assertEqual(table[3].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(
          _get_shapes(table[2].inputs),
          ((batch_size, 28, 28, 32), {"training": True}),
      )
      self.assertEqual(table[2].outputs.shape, (batch_size, 28, 28, 64))
  
      self.assertEqual(table[3].inputs.shape, (batch_size, 64))
      self.assertEqual(table[3].outputs.shape, (batch_size, 10))
  
      # check ConvBlock_0 and ConvBlock_1 are summarized
      self.assertNotEqual(
          table[1].module_variables, table[1].counted_variables
      )
      self.assertNotEqual(
          table[2].module_variables, table[2].counted_variables
      )
  
      # check CNN and Dense_0 output are not summarized
      self.assertEqual(
          table[0].module_variables, table[0].counted_variables
      )
      self.assertEqual(
          table[3].module_variables, table[3].counted_variables
      )
  
      # Top level FLOPs > sum of listed submodule FLOPs, since not all are listed.
      self.assertGreater(
          table[0].flops, sum(r.flops for r in table[1:])
      )
      self.assertGreater(
          table[0].vjp_flops, sum(r.vjp_flops for r in table[1:])
      )
  
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 2 for calling the SummaryTest obj's underlying member method code with example usage: obj.test_module_summary_with_depth() without return types