code: def test_module_reuse(self):
      class ConvBlock(nn.Module):
          @nn.compact
          def __call__(self, x):
              x = nn.Conv(32, [3, 3])(x)
              x = nn.BatchNorm(use_running_average=True)(x)
              x = nn.Dropout(0.5, deterministic=True)(x)
              x = nn.relu(x)
              return x
  
      class CNN(nn.Module):
          @nn.compact
          def __call__(self, x):
              block = ConvBlock()
              x = block(x)
              x = block(x)
              x = block(x)
              return x
  
      x = jnp.ones((4, 28, 28, 32))
      module_repr = CNN().tabulate(
          jax.random.key(0),
          x=x,
          show_repeated=True,
          console_kwargs=CONSOLE_TEST_KWARGS,
          compute_flops=True,
          compute_vjp_flops=True,
      )
      lines = module_repr.splitlines()
  
      # first call
      self.assertIn("ConvBlock_0/Conv_0", lines[9])
      self.assertIn("bias", lines[9])
      self.assertIn("ConvBlock_0/BatchNorm_0", lines[14])
      self.assertIn("mean", lines[14])
      self.assertIn("bias", lines[14])
      self.assertIn("ConvBlock_0/Dropout_0", lines[19])
  
      # second call
      self.assertIn("ConvBlock_0/Conv_0", lines[23])
      self.assertNotIn("bias", lines[23])
      self.assertIn("ConvBlock_0/BatchNorm_0", lines[25])
      self.assertNotIn("mean", lines[25])
      self.assertNotIn("bias", lines[25])
      self.assertIn("ConvBlock_0/Dropout_0", lines[27])
  
      # third call
      self.assertIn("ConvBlock_0/Conv_0", lines[31])
      self.assertNotIn("bias", lines[31])
      self.assertIn("ConvBlock_0/BatchNorm_0", lines[33])
      self.assertNotIn("mean", lines[33])
      self.assertNotIn("bias", lines[33])
      self.assertIn("ConvBlock_0/Dropout_0", lines[35])
  
      # Test that CNN FLOPs are 3x ConvBlock FLOPs.
      args = ({"dropout": random.key(0), "params": random.key(1)}, x)
      cnn = summary._get_module_table(
          CNN(),
          depth=1,
          show_repeated=True,
          compute_flops=True,
          compute_vjp_flops=True,
      )(*args, mutable=True)
  
      block = summary._get_module_table(
          ConvBlock(),
          depth=1,
          show_repeated=True,
          compute_flops=True,
          compute_vjp_flops=True,
      )(*args, mutable=True)
  
      # Total forward/backward FLOPs equal to their sums of sub blocks.
      self.assertEqual(cnn[0].flops, sum(r.flops for r in cnn[1:]))
      self.assertEqual(
          cnn[0].vjp_flops, sum(r.vjp_flops for r in cnn[1:])
      )
  
      # Each sub block has cost equal to ConvBlock instantiated separately.
      for r in cnn[1:]:
          self.assertEqual(r.flops, block[0].flops)
          self.assertEqual(r.vjp_flops, block[0].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 9 for calling the SummaryTest obj's underlying member method code with example usage: obj.test_module_reuse() and returns: x, x