code: def test_partitioned_params(self):
      class Classifier(nn.Module):
          @nn.compact
          def __call__(self, x):
              hidden = nn.Dense(
                  features=1024,
                  kernel_init=nn.with_partitioning(
                      nn.initializers.lecun_normal(), (None, "data")
                  ),
                  bias_init=nn.with_partitioning(
                      nn.initializers.zeros, (None,)
                  ),
                  name="hidden",
              )
              x = x / 255.0
              x = x.reshape((x.shape[0], -1))  # flatten
              x = nn.relu(hidden(x))
              x = nn.Dense(features=10, name="head")(x)
              return x
  
      module = Classifier()
      lines = module.tabulate(
          jax.random.key(0),
          jnp.empty((1, 28, 28, 1)),
          console_kwargs=CONSOLE_TEST_KWARGS,
          compute_flops=True,
          compute_vjp_flops=True,
      ).splitlines()
      self.assertIn("P(None,)", lines[7])
      self.assertIn("P(None, data)", lines[8])
  
      # Per-layer forward FLOPs:
      self.assertIn(
          "1606656", lines[7]
      )  # 1 * (28 * 28 * 1) * 1024 * 2 + 1024
      self.assertIn(
          "20490", lines[12]
      )  #  1 * (1024       ) * 10   * 2 + 10
  
      # Total forward FLOPs: input division + ReLU + two dense layers above.
      # (1 * 28 * 28 * 1) + (1 * 1024) + 1606656 + 20490.
      self.assertIn("1628954", lines[5])
  
      # Per-layer backward FLOPs:
  
      # [3x MMs: forward, input cotangent, weight cotangent] 1024 * 784 * 2 * 3
      # + [forward bias addition] 1024
      # + [`mutable=True`: weight and bias sizes] 1024 * 784 + 1024
      self.assertIn("5621760", lines[7])
  
      # [3x matmuls: forward, input cotangent, weight cotangent] 1024 * 10 * 2 * 3
      # + [forward bias addition] 10
      # + [`mutable=True`: weight and bias sizes] 1024 * 10 + 10
      self.assertIn("71700", lines[12])
  
      # Total backward FLOPs: input division + ReLU + two dense layers above.
      # 2 * (1 * 28 * 28 * 1) + 3 * (1 * 1024) + 5621760 + 71700.
      self.assertIn("5698100", lines[5])
  
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 12 for calling the SummaryTest obj's underlying member method code with example usage: obj.test_partitioned_params() and returns: x