Repository: google/flax

Issue:
`nn.tabulate` results in `KeyError: 'intermediates'` with methods that include transformations
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Arch Linux x64
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: `flax=0.5.1`, `jax=0.3.13`, `jaxlib=0.3.10`
- Python version: 3.10.5
- GPU/TPU model and memory: RTX 3080 12GB
- CUDA version (if applicable): 11.7

### Problem you have encountered:

`nn.tabulate` results in KeyError: 'intermediates' when used with methods that include transformations. Tested with `vmap` and `scan` (included repro below).

### Logs, error messages, etc:
```
File ~/.local/lib/python3.10/site-packages/flax/linen/summary.py:173, in tabulate.<locals>._tabulate_fn(*args, **kwargs)
    169 def _tabulate_fn(*args, **kwargs):
    170   table_fn = _get_module_table(module, rngs, method=method, 
    171                                mutable=mutable, depth=depth, 
    172                                exclude_methods=set(exclude_methods))
--> 173   table = table_fn(*args, **kwargs)
    174   return _render_table(table)

File ~/.local/lib/python3.10/site-packages/flax/linen/summary.py:213, in _get_module_table.<locals>._get_table_fn(*args, **kwargs)
    209 collections: List[str] = [
    210     col for col in shape_variables.keys() if col != 'intermediates'
    211 ]
    212 shape_variables = shape_variables.unfreeze()
--> 213 rows = list(
    214     _flatten_to_rows(
    215         path=(),
    216         variables=shape_variables,
    217         depth=depth,
...
    250   }
    252   if len(module_outputs) == 0:
    253     output = None

KeyError: 'intermediates'
```


### Steps to reproduce:

Minimal repro:

```python
import jax.numpy as jnp
from jax import random
import flax.linen as nn


class LSTM(nn.Module):
    batch_size: int
    out_feat: int

    @nn.compact
    def __call__(self, x):
        carry = nn.LSTMCell.initialize_carry(
            random.PRNGKey(0), (self.batch_size,), self.out_feat
        )
        Cell = nn.scan(
            nn.LSTMCell,
            variable_broadcast="params",
            split_rngs={"params": False},
            in_axes=1,
            out_axes=1,
        )
        return Cell()(carry, x)


if __name__ == "__main__":
    lstm = LSTM(batch_size=128, out_feat=128)

    # KeyError: 'intermediates'
    print(lstm.tabulate(random.PRNGKey(0), jnp.ones((128, 128))))
```


Discussion:
Hey @RocketLL, thanks for minimal repro. 

@jheek @marcvanzee  the `Cell` Module (`ScanLSTMCell_0`) is neither reporting its outputs nor the output of its submodules as shown here:

```
{
    intermediates: {
        __call__: (((ShapeDtypeStruct(shape=(128, 128), dtype=float32), ShapeDtypeStruct(shape=(128, 128), dtype=float32)), ShapeDtypeStruct(shape=(128, 128, 128), dtype=float32)),),
    },
    params: {
        ScanLSTMCell_0: {
            hf: {
                bias: ShapeDtypeStruct(shape=(128,), dtype=float32),
                kernel: ShapeDtypeStruct(shape=(128, 128), dtype=float32),
            },
            hg: {
                bias: ShapeDtypeStruct(shape=(128,), dtype=float32),
                kernel: ShapeDtypeStruct(shape=(128, 128), dtype=float32),
            },
            hi: {
                bias: ShapeDtypeStruct(shape=(128,), dtype=float32),
                kernel: ShapeDtypeStruct(shape=(128, 128), dtype=float32),
            },
            ho: {
                bias: ShapeDtypeStruct(shape=(128,), dtype=float32),
                kernel: ShapeDtypeStruct(shape=(128, 128), dtype=float32),
            },
            if: {
                kernel: ShapeDtypeStruct(shape=(128, 128), dtype=float32),
            },
            ig: {
                kernel: ShapeDtypeStruct(shape=(128, 128), dtype=float32),
            },
            ii: {
                kernel: ShapeDtypeStruct(shape=(128, 128), dtype=float32),
            },
            io: {
                kernel: ShapeDtypeStruct(shape=(128, 128), dtype=float32),
            },
        },
    },
}
```
Is this expected? 

For now I'll have to rethink the policy of how to detect submodules, there are two options:

1. Treat `ScanLSTMCell_0`, `hf`, etc, as submodules with no outputs.
2. Treat the whole `ScanLSTMCell_0` structure as just parameters.
@jheek points out that this can be fixed by adding `intermediates` to `variable_axes`:

```python
Cell = nn.scan(
    nn.LSTMCell,
    variable_broadcast="params",
    split_rngs={"params": False},
    variable_axes={"intermediates": 1},
    in_axes=1,
    out_axes=1,
)
```
However the internal use of `capture_intermediates` is an implementation detail and the users shouldn't be required to be aware of this. I'll create a new issue to track this broader problem with `tabulate` and propose a more general solution. 

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -18,17 +18,20 @@
 import enum
 import functools
 import inspect
+import re
 import threading
 import typing
 import weakref
-from typing import (Any, Callable, Dict, Generic, Iterable, List, Optional,
-                    Sequence, Set, Tuple, Type, TypeVar, Union, overload)
+from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, Mapping,
+                    Optional, Set, Tuple, Type, TypeVar, Union, overload)
 
 import jax
 import numpy as np
+import jax.numpy as jnp
 from typing_extensions import \
     dataclass_transform  # pytype: disable=not-supported-yet
 
+import flax
 from flax import (config, core, errors, serialization, traceback_util,
                   traverse_util)
 from flax.core import Scope
@@ -37,8 +40,6 @@
     CollectionFilter, DenyList, FrozenVariableDict, Variable, VariableDict,
     union_filters)
 from flax.ids import uuid
-from flax.linen import summary
-
 
 
 traceback_util.register_exclusion(__file__)
@@ -61,6 +62,16 @@
 
 # pylint: disable=protected-access,attribute-defined-outside-init
 
+def _get_value_representation(x: Any) -> 'flax.linen.summary._ValueRepresentation':
+  from flax.linen import summary
+
+  if isinstance(x, (int, float, bool, type(None))) or (
+    isinstance(x, np.ndarray) and np.isscalar(x)):
+    return summary._ObjectRepresentation(x)
+  try:
+    return summary._ArrayRepresentation(jnp.shape(x), jnp.result_type(x))
+  except:
+    return summary._ObjectRepresentation(x)
 
 def _indent(x: str, num_spaces: int):
   indent_str = ' ' * num_spaces
@@ -104,6 +115,46 @@ def _module_repr(module: 'Module', num_spaces: int = 4):
   else:
     return f'{cls_name}()'
 
+# 
+# -----------------------------------------------------------------------------
+
+_find_non_lifted_module = re.compile(r'.*\((.*)\)')
+
+def _fix_path_part(part: str):
+  """Fixes a path part by removing transformation name and parenthesis sometimes 
+  inserted by lifted transformations"""
+  match = _find_non_lifted_module.match(part)
+  if match:
+    return match.group(1)
+  return part
+
+@dataclasses.dataclass
+class _CallInfo:
+  index: int
+  path: Tuple[str, ...]
+  module_type: Type['Module']
+  method: str
+  args: Tuple[Any, ...]
+  kwargs: Dict[str, Any]
+  outputs: Any
+
+@dataclasses.dataclass
+class _CallInfoContext(threading.local):
+  index: int
+  calls: List[_CallInfo]
+
+  def get_call_index(self, module: 'Module') -> int:
+    index = self.index
+    self.index += 1
+    return index
+
+@contextlib.contextmanager
+def _tabulate_context():
+  _context.call_info_stack.append(_CallInfoContext(0, []))
+  try:
+    yield
+  finally:
+    _context.call_info_stack.pop()
 
 # Track parent relationship across Modules.
 # -----------------------------------------------------------------------------
@@ -128,6 +179,13 @@ def capture_stack(self):
       self._thread_data.capture_stack = []
     return self._thread_data.capture_stack
 
+  @property
+  def call_info_stack(self) -> List[_CallInfoContext]:
+    """Keeps track of the active call_info_context."""
+    if not hasattr(self._thread_data, 'call_info_stack'):
+      self._thread_data.call_info_stack = []
+    return self._thread_data.call_info_stack
+
 # The global context
 _context = _DynamicContext()
 
@@ -638,6 +696,7 @@ def _call_wrapped_method(self, fun, args, kwargs):
     is_compact_method = hasattr(fun, 'compact')
     fun_name = getattr(fun, '__name__', 'unnamed_function')
     is_setup_method = fun_name == 'setup'
+    add_call_info = not is_setup_method and len(_context.call_info_stack) > 0
     # We lazily call setup() only when needed.
     if is_setup_method:
       is_recurrent = self._state.in_setup
@@ -652,15 +711,27 @@ def _call_wrapped_method(self, fun, args, kwargs):
       self._state.in_compact_method = True
     _context.module_stack.append(self)
     try:
+      # get call info
+      if add_call_info:
+        call_index = _context.call_info_stack[-1].get_call_index(self)
+        scope_path = jax.tree_util.tree_map(_fix_path_part, self.scope.path)
+      
+      # call method
       if _use_named_call:
         with jax.named_scope(_derive_profiling_name(self, fun)):
           y = fun(self, *args, **kwargs)
       else:
         y = fun(self, *args, **kwargs)
+
       if _context.capture_stack:
         filter_fn = _context.capture_stack[-1]
         if filter_fn and filter_fn(self, fun_name):
           self.sow('intermediates', fun_name, y)
+      if add_call_info:
+        _args, _kwargs, _y = jax.tree_util.tree_map(
+          _get_value_representation, (args, kwargs, y), is_leaf=lambda x: x is None)
+        _context.call_info_stack[-1].calls.append(
+          _CallInfo(call_index, scope_path, type(self), fun.__name__, _args, _kwargs, _y))
       return y
     finally:
       _context.module_stack.pop()
@@ -1410,17 +1481,17 @@ def tabulate(
     self,
     rngs: Union[PRNGKey, RNGSequences],
     *args,
-    method: Optional[Callable[..., Any]] = None,
-    mutable: CollectionFilter = True,
     depth: Optional[int] = None,
-    exclude_methods: Sequence[str] = (),
+    show_repeated: bool = False,
+    mutable: CollectionFilter = True,
+    console_kwargs: Optional[Mapping[str, Any]] = None,
     **kwargs) -> str:
     """Creates a summary of the Module represented as a table.
 
-    This method has the same signature as `init`, but instead of returning
-    the variables, it returns the string summarizing the Module in a table.
-    `tabulate` uses `jax.eval_shape` to run the forward computation without
-    consuming any FLOPs or allocating memory.
+    This method has the same signature and internally calls `Module.init`, 
+    but instead of returning the variables, it returns the string summarizing 
+    the Module in a table. `tabulate` uses `jax.eval_shape` to run the forward 
+    computation without consuming any FLOPs or allocating memory.
 
     Example::
 
@@ -1441,61 +1512,60 @@ def __call__(self, x):
 
     This gives the following output::
 
-                         Foo Summary
-      ┏━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
-      ┃ path    ┃ outputs       ┃ params               ┃
-      ┡━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
-      │ Inputs  │ float32[16,9] │                      │
-      ├─────────┼───────────────┼──────────────────────┤
-      │ Dense_0 │ float32[16,4] │ bias: float32[4]     │
-      │         │               │ kernel: float32[9,4] │
-      │         │               │                      │
-      │         │               │ 40 (160 B)           │
-      ├─────────┼───────────────┼──────────────────────┤
-      │ Dense_1 │ float32[16,2] │ bias: float32[2]     │
-      │         │               │ kernel: float32[4,2] │
-      │         │               │                      │
-      │         │               │ 10 (40 B)            │
-      ├─────────┼───────────────┼──────────────────────┤
-      │ Foo     │ float32[16,2] │                      │
-      ├─────────┼───────────────┼──────────────────────┤
-      │         │         Total │ 50 (200 B)           │
-      └─────────┴───────────────┴──────────────────────┘
-
-                Total Parameters: 50 (200 B)
+                                      Foo Summary                                
+      ┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
+      ┃ path    ┃ module ┃ inputs        ┃ outputs       ┃ params               ┃
+      ┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
+      │         │ Foo    │ float32[16,9] │ float32[16,2] │                      │
+      ├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
+      │ Dense_0 │ Dense  │ float32[16,9] │ float32[16,4] │ bias: float32[4]     │
+      │         │        │               │               │ kernel: float32[9,4] │
+      │         │        │               │               │                      │
+      │         │        │               │               │ 40 (160 B)           │
+      ├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
+      │ Dense_1 │ Dense  │ float32[16,4] │ float32[16,2] │ bias: float32[2]     │
+      │         │        │               │               │ kernel: float32[4,2] │
+      │         │        │               │               │                      │
+      │         │        │               │               │ 10 (40 B)            │
+      ├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
+      │         │        │               │         Total │ 50 (200 B)           │
+      └─────────┴────────┴───────────────┴───────────────┴──────────────────────┘
+                                                                                
+                            Total Parameters: 50 (200 B)
 
     **Note**: rows order in the table does not represent execution order,
     instead it aligns with the order of keys in `variables` which are sorted
     alphabetically.
 
     Args:
-      rngs: The rngs for the variable collections.
+      rngs: The rngs for the variable collections as passed to `Module.init`.
       *args: The arguments to the forward computation.
-      method: An optional method. If provided, applies this method. If not
-        provided, applies the ``__call__`` method.
-      mutable: Can be bool, str, or list. Specifies which collections should be
-        treated as mutable: ``bool``: all/no collections are mutable.
-        ``str``: The name of a single mutable collection. ``list``: A
-        list of names of mutable collections. By default all collections
-        except 'intermediates' are mutable.
       depth: controls how many submodule deep the summary can go. By default its
         `None` which means no limit. If a submodule is not shown because of the
-        depth limit, its parameter count and bytes will be added to the row of
-        its first shown ancestor such that the sum of all rows always adds up to
-        the total number of parameters of the Module.
-      exclude_methods: A sequence of strings that specifies which methods should
-        be ignored. In case a module calls a helper method from its main method,
-        use this argument to exclude the helper method from the summary to avoid
-        ambiguity.
+        depth limit, its parameter count and bytes will be added to the row of its
+        first shown ancestor such that the sum of all rows always adds up to the
+        total number of parameters of the Module.
+      show_repeated: If `True`, repeated calls to the same module will be shown
+        in the table, otherwise only the first call will be shown. Default is
+        `False`.
+      mutable: Can be bool, str, or list. Specifies which collections should be
+        treated as mutable: ``bool``: all/no collections are mutable. ``str``: The
+        name of a single mutable collection. ``list``: A list of names of mutable
+        collections. By default all collections except 'intermediates' are
+        mutable.
+      console_kwargs: An optional dictionary with additional keyword arguments that 
+        are passed to `rich.console.Console` when rendering the table. Default arguments 
+        are `{'force_terminal': True, 'force_jupyter': False}`.
       **kwargs: keyword arguments to pass to the forward computation.
 
     Returns:
       A string summarizing the Module.
     """
-
-    tabulate_fn = summary.tabulate(self, rngs, method=method,
-                                   mutable=mutable, depth=depth,
-                                   exclude_methods=exclude_methods)
+    from flax.linen import summary
+    
+    tabulate_fn = summary.tabulate(self, rngs, depth=depth,
+                                   show_repeated=show_repeated, mutable=mutable,
+                                   console_kwargs=console_kwargs)
     return tabulate_fn(*args, **kwargs)
 
 
diff --git a/flax/linen/summary.py b/flax/linen/summary.py
--- a/flax/linen/summary.py
+++ b/flax/linen/summary.py
@@ -13,13 +13,15 @@
 # limitations under the License.
 
 """Flax Module summary library."""
+from abc import ABC, abstractmethod
 import dataclasses
 import io
-from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union
 
-import flax
-from flax.core.scope import CollectionFilter, DenyList
+import flax.linen.module as module_lib
+from flax.core.scope import CollectionFilter, FrozenVariableDict, MutableVariableDict
 import jax
+import jax.numpy as jnp
 import rich.console
 import rich.table
 import rich.text
@@ -29,6 +31,42 @@
 RNGSequences = Dict[str, PRNGKey]
 Array = Any    # pylint: disable=invalid-name
 
+class _ValueRepresentation(ABC):
+  """A class that represents a value in the summary table."""
+
+  @abstractmethod
+  def render(self) -> str:
+    ...
+
+  @abstractmethod
+  def value(self) -> Any:
+    ...
+
+@dataclasses.dataclass
+class _ArrayRepresentation(_ValueRepresentation):
+  shape: Tuple[int, ...]
+  dtype: Any
+
+  @classmethod
+  def render_array(cls, x) -> str:
+    return cls(jnp.shape(x), jnp.result_type(x)).render()
+
+  def render(self):
+    shape_repr = ','.join(str(x) for x in self.shape)
+    return f'[dim]{self.dtype}[/dim][{shape_repr}]'
+
+  def value(self):
+    return self
+
+@dataclasses.dataclass
+class _ObjectRepresentation(_ValueRepresentation):
+  obj: Any
+
+  def render(self):
+    return repr(self.obj)
+
+  def value(self):
+    return self.obj
 
 @dataclasses.dataclass
 class Row:
@@ -46,12 +84,18 @@ class Row:
       from submodules depending on the depth of the Module in question.
   """
   path: Tuple[str, ...]
+  module_type: Type[module_lib.Module]
+  method: str
+  inputs: Any
   outputs: Any
-  module_variables: Dict[str, Dict[str, Array]]
-  counted_variables: Dict[str, Dict[str, Array]]
+  module_variables: Dict[str, Dict[str, Any]]
+  counted_variables: Dict[str, Dict[str, Any]]
+
+  def __post_init__(self):
+    self.inputs = _normalize_structure(self.inputs)
+    self.outputs = _normalize_structure(self.outputs)
 
-  def size_and_bytes(self,
-                     collections: Iterable[str]) -> Dict[str, Tuple[int, int]]:
+  def size_and_bytes(self, collections: Iterable[str]) -> Dict[str, Tuple[int, int]]:
     return {
         col: _size_and_bytes(self.counted_variables[col])
         if col in self.counted_variables else (0, 0) for col in collections
@@ -68,7 +112,7 @@ class Table(List[Row]):
   * `collections`: a list containing the parameter collections (e.g. 'params', 'batch_stats', etc)
   """
 
-  def __init__(self, module: 'flax.linen.Module', collections: List[str],
+  def __init__(self, module: module_lib.Module, collections: Sequence[str],
                rows: Iterable[Row]):
     super().__init__(rows)
     self.module = module
@@ -76,22 +120,21 @@ def __init__(self, module: 'flax.linen.Module', collections: List[str],
 
 
 def tabulate(
-    module: 'flax.linen.Module',
-    rngs: Union[PRNGKey, RNGSequences],
-    method: Optional[Callable[..., Any]] = None,
-    mutable: CollectionFilter = True,
-    depth: Optional[int] = None,
-    exclude_methods: Sequence[str] = (),
+  module: module_lib.Module,
+  rngs: Union[PRNGKey, RNGSequences],
+  depth: Optional[int] = None,
+  show_repeated: bool = False,
+  mutable: CollectionFilter = True,
+  console_kwargs: Optional[Mapping[str, Any]] = None,
+  **kwargs,
 ) -> Callable[..., str]:
   """Returns a function that creates a summary of the Module represented as a table.
 
-  This function accepts most of the same arguments as `Module.init`, except that
-  it returns a function of the form `(*args, **kwargs) -> str` where `*args` and
-  `**kwargs`
-  are passed to `method` (e.g. `__call__`) during the forward pass.
+  This function accepts most of the same arguments and internally calls `Module.init`, 
+  except that it returns a function of the form `(*args, **kwargs) -> str` where `*args` 
+  and `**kwargs` are passed to `method` (e.g. `__call__`) during the forward pass.
 
-  `tabulate` uses `jax.eval_shape` under the hood to run the forward computation
-  without
+  `tabulate` uses `jax.eval_shape` under the hood to run the forward computation without
   consuming any FLOPs or allocating memory.
 
   Example::
@@ -101,10 +144,10 @@ def tabulate(
     import flax.linen as nn
 
     class Foo(nn.Module):
-        @nn.compact
-        def __call__(self, x):
-            h = nn.Dense(4)(x)
-            return nn.Dense(2)(h)
+      @nn.compact
+      def __call__(self, x):
+        h = nn.Dense(4)(x)
+        return nn.Dense(2)(h)
 
     x = jnp.ones((16, 9))
     tabulate_fn = nn.tabulate(Foo(), jax.random.PRNGKey(0))
@@ -114,28 +157,27 @@ def __call__(self, x):
 
   This gives the following output::
 
-                          Foo Summary
-      ┏━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
-      ┃ path    ┃ outputs       ┃ params               ┃
-      ┡━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
-      │ Inputs  │ float32[16,9] │                      │
-      ├─────────┼───────────────┼──────────────────────┤
-      │ Dense_0 │ float32[16,4] │ bias: float32[4]     │
-      │         │               │ kernel: float32[9,4] │
-      │         │               │                      │
-      │         │               │ 40 (160 B)           │
-      ├─────────┼───────────────┼──────────────────────┤
-      │ Dense_1 │ float32[16,2] │ bias: float32[2]     │
-      │         │               │ kernel: float32[4,2] │
-      │         │               │                      │
-      │         │               │ 10 (40 B)            │
-      ├─────────┼───────────────┼──────────────────────┤
-      │ Foo     │ float32[16,2] │                      │
-      ├─────────┼───────────────┼──────────────────────┤
-      │         │         Total │ 50 (200 B)           │
-      └─────────┴───────────────┴──────────────────────┘
-
-                  Total Parameters: 50 (200 B)
+                                    Foo Summary                                
+    ┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
+    ┃ path    ┃ module ┃ inputs        ┃ outputs       ┃ params               ┃
+    ┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
+    │         │ Foo    │ float32[16,9] │ float32[16,2] │                      │
+    ├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
+    │ Dense_0 │ Dense  │ float32[16,9] │ float32[16,4] │ bias: float32[4]     │
+    │         │        │               │               │ kernel: float32[9,4] │
+    │         │        │               │               │                      │
+    │         │        │               │               │ 40 (160 B)           │
+    ├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
+    │ Dense_1 │ Dense  │ float32[16,4] │ float32[16,2] │ bias: float32[2]     │
+    │         │        │               │               │ kernel: float32[4,2] │
+    │         │        │               │               │                      │
+    │         │        │               │               │ 10 (40 B)            │
+    ├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
+    │         │        │               │         Total │ 50 (200 B)           │
+    └─────────┴────────┴───────────────┴───────────────┴──────────────────────┘
+                                                                              
+                          Total Parameters: 50 (200 B)                        
+
 
   **Note**: rows order in the table does not represent execution order,
   instead it aligns with the order of keys in `variables` which are sorted
@@ -143,22 +185,24 @@ def __call__(self, x):
 
   Args:
     module: The module to tabulate.
-    method: An optional method. If provided, applies this method. If not
-      provided, applies the ``__call__`` method.
-    mutable: Can be bool, str, or list. Specifies which collections should be
-      treated as mutable: ``bool``: all/no collections are mutable. ``str``: The
-      name of a single mutable collection. ``list``: A list of names of mutable
-      collections. By default all collections except 'intermediates' are
-      mutable.
+    rngs: The rngs for the variable collections as passed to `Module.init`.
     depth: controls how many submodule deep the summary can go. By default its
       `None` which means no limit. If a submodule is not shown because of the
       depth limit, its parameter count and bytes will be added to the row of its
       first shown ancestor such that the sum of all rows always adds up to the
       total number of parameters of the Module.
-    exclude_methods: A sequence of strings that specifies which methods should
-      be ignored. In case a module calls a helper method from its main method,
-      use this argument to exclude the helper method from the summary to avoid
-      ambiguity.
+    mutable: Can be bool, str, or list. Specifies which collections should be
+      treated as mutable: ``bool``: all/no collections are mutable. ``str``: The
+      name of a single mutable collection. ``list``: A list of names of mutable
+      collections. By default all collections except 'intermediates' are
+      mutable.
+    show_repeated: If `True`, repeated calls to the same module will be shown
+      in the table, otherwise only the first call will be shown. Default is
+      `False`.
+    console_kwargs: An optional dictionary with additional keyword arguments that 
+      are passed to `rich.console.Console` when rendering the table. Default arguments 
+      are `{'force_terminal': True, 'force_jupyter': False}`.
+    **kwargs: Additional arguments passed to `Module.init`.
 
   Returns:
     A function that accepts the same `*args` and `**kwargs` of the forward pass
@@ -166,170 +210,125 @@ def __call__(self, x):
     Modules.
   """
 
-  def _tabulate_fn(*args, **kwargs):
-    table_fn = _get_module_table(module, rngs, method=method, 
-                                 mutable=mutable, depth=depth, 
-                                 exclude_methods=set(exclude_methods))
-    table = table_fn(*args, **kwargs)
-    return _render_table(table)
+  def _tabulate_fn(*fn_args, **fn_kwargs):
+    table_fn = _get_module_table(module, depth=depth, show_repeated=show_repeated)
+    table = table_fn(rngs, *fn_args, mutable=mutable, **fn_kwargs, **kwargs)
+    return _render_table(table, console_kwargs)
 
   return _tabulate_fn
 
-
 def _get_module_table(
-    module: 'flax.linen.Module',
-    rngs: Union[PRNGKey, RNGSequences],
-    method: Optional[Callable[..., Any]],
-    mutable: CollectionFilter,
+    module: module_lib.Module,
     depth: Optional[int],
-    exclude_methods: Set[str],
+    show_repeated: bool,
 ) -> Callable[..., Table]:
-
-  exclude_methods.add("setup")
+  """A function that takes a Module and returns function with the same signature as `init`
+  but returns the Table representation of the Module."""
 
   def _get_table_fn(*args, **kwargs):
-    output_methods: Set[str] = set()
-
-    def capture_intermediates(_module, method_name: str):
-      if method_name in exclude_methods:
-        return False
+    
+    with module_lib._tabulate_context():
+
+      def _get_variables():
+        return module.init(*args, **kwargs)
+
+      variables = jax.eval_shape(_get_variables)
+      calls = module_lib._context.call_info_stack[-1].calls
+      calls.sort(key=lambda c: c.index)
+
+    collections: Set[str] = set(variables.keys())
+    rows = []
+    all_paths: Set[Tuple[str, ...]] = set(call.path for call in calls)
+    visited_paths: Set[Tuple[str, ...]] = set()
+
+    for c in calls:
+      call_depth = len(c.path)
+      inputs = _process_inputs(c.args, c.kwargs)
+
+      if c.path in visited_paths:
+        if not show_repeated:
+          continue
+        module_vars = {}
+        counted_vars = {}
+      elif depth is not None:
+        if call_depth > depth:
+          continue
+        module_vars, _ = _get_module_variables(c.path, variables, all_paths)
+        if call_depth == depth:
+          counted_vars = _get_path_variables(c.path, variables)
+        else:
+          counted_vars = module_vars
       else:
-        output_methods.add(method_name)
-        return True
-
-    shape_variables = jax.eval_shape(lambda: module.init(
-        rngs,
-        *args,
-        method=method,
-        mutable=mutable,
-        capture_intermediates=capture_intermediates,
-        **kwargs,
-    ))
-
-    collections: List[str] = [
-        col for col in shape_variables.keys() if col != 'intermediates'
-    ]
-    shape_variables = shape_variables.unfreeze()
-    rows = list(
-        _flatten_to_rows(
-            path=(),
-            variables=shape_variables,
-            depth=depth,
-            output_methods=output_methods))
-
-    if args and kwargs:
-      input_values = (*args, kwargs)
-    elif args and not kwargs:
-      input_values = args[0] if len(args) == 1 else args
-    elif kwargs and not args:
-      input_values = kwargs
-    else:
-      input_values = ''
-
-    inputs_row = Row(('Inputs',), input_values, {}, {})
-    rows.insert(0, inputs_row)
-
-    return Table(module, collections, rows)
+        module_vars, _ = _get_module_variables(c.path, variables, all_paths)
+        counted_vars = module_vars
+        
+      visited_paths.add(c.path)
+      rows.append(
+        Row(c.path, c.module_type, c.method, inputs, c.outputs, module_vars, counted_vars))
+        
+    return Table(module, tuple(collections), rows)
 
   return _get_table_fn
 
-
-def _flatten_to_rows(
-    path: Tuple[str, ...],
-    variables: Dict[str, Any],
-    depth: Optional[int],
-    output_methods: Set[str],
-) -> Iterable[Row]:
-
-  # get variables only for this Module
-  module_variables = _get_module_variables(variables)
-  module_outputs = {
-      key: value
-      for key, value in variables['intermediates'].items()
-      if key in output_methods
-  }
-
-  if len(module_outputs) == 0:
-    output = None
-  elif len(module_outputs) > 1:
-    raise ValueError(
-        f"Cannot infer output, module '{'/'.join(path)}' has multiple "
-        f"intermediates: {list(module_outputs.keys())}. Use the `exclude_methods` "
-        f"argument to make sure each module only reports one output.")
-  else:
-    output = list(module_outputs.values())[0][0]
-
-  if depth is not None and depth == 0:
-    # don't recurse, yield current level
-    # count_variables contains all variables that are not intermediates
-    variables = variables.copy()
-    del variables['intermediates']
-    module_variables.pop('intermediates')
-    yield Row(
-        path=path,
-        outputs=output,
-        module_variables=module_variables,
-        counted_variables=variables,
-    )
+def _get_module_variables(
+  path: Tuple[str, ...], variables: FrozenVariableDict, all_paths: Set[Tuple[str, ...]]
+) -> Tuple[MutableVariableDict, Any]:
+  """A function that takes a path and variables structure and returns a 
+  (module_variables, submodule_variables) tuple for that path. _get_module_variables
+  uses the `all_paths` set to determine if a variable belongs to a submodule or not."""
+  module_variables = _get_path_variables(path, variables)
+  submodule_variables = {collection: {} for collection in module_variables}
+  all_keys = set(key for collection in module_variables.values() for key in collection)
+
+  for key in all_keys:
+    submodule_path = path + (key,)
+    if submodule_path in all_paths:
+
+      for collection in module_variables:
+        if key in module_variables[collection]:
+          submodule_variables[collection][key] = module_variables[collection].pop(key)
+
+  return module_variables, submodule_variables
+
+def _get_path_variables(path: Tuple[str, ...], variables: FrozenVariableDict) -> MutableVariableDict:
+  """A function that takes a path and a variables structure and returns the variable structure at 
+  that path."""
+  path_variables = {}
+
+  for collection in variables:
+    collection_variables = variables[collection]
+    for name in path:
+      if name not in collection_variables:
+        collection_variables = None
+        break
+      collection_variables = collection_variables[name]
+
+    if collection_variables is not None:
+      path_variables[collection] = collection_variables.unfreeze()
+    
+  return path_variables
+
+def _process_inputs(args, kwargs) -> Any:
+  """A function that normalizes the representation of the ``args`` and ``kwargs``
+  for the ``inputs`` column."""
+  if args and kwargs:
+    input_values = (*args, kwargs)
+  elif args and not kwargs:
+    input_values = args[0] if len(args) == 1 else args
+  elif kwargs and not args:
+    input_values = kwargs
   else:
-    # recurse into lower levels
-    keys = list(key for key in variables['intermediates'].keys()
-                if key not in module_variables['intermediates'])
-
-    # add keys from other collections
-    # dont use set here because we want to preserve order
-    for collection in variables:
-      if collection != 'intermediates':
-        for key in variables[collection]:
-          if key not in keys and key not in module_variables.get(
-              collection, {}):
-            keys.append(key)
-
-    for key in keys:
-      next_path = path + (key,)
-      next_variables = _step_into(variables, key)
-      yield from _flatten_to_rows(
-          path=next_path,
-          variables=next_variables,
-          depth=depth - 1 if depth is not None else None,
-          output_methods=output_methods,
-      )
-
-    # current row
-    yield Row(
-        path=path,
-        outputs=output,
-        module_variables=module_variables,
-        counted_variables=module_variables,
-    )
-
-
-def _step_into(variables: Dict[str, Any], key: str):
-  return {
-      col: params[key] for col, params in variables.items() if key in params
-  }
+    input_values = ()
 
+  return input_values
 
-def _get_module_variables(variables: Dict[str, Any]) -> Dict[str, Any]:
-
-  module_variables: Dict[str, Dict[str, Any]] = {
-      collection: {
-          name: value
-          for name, value in params.items()
-          if not isinstance(value, Mapping)  # is this robust?
-      } for collection, params in variables.items()
-  }
-  # filter empty collectionswhen
-  module_variables = {
-      collection: params
-      for collection, params in module_variables.items()
-      if len(params) > 0
-  }
-
-  return module_variables
-
-
-def _render_table(table: Table) -> str:
+def _render_table(table: Table, console_extras: Optional[Mapping[str, Any]]) -> str:
+  """A function that renders a Table to a string representation using rich."""
+  console_kwargs = {'force_terminal': True, 'force_jupyter': False}
+  if console_extras is not None:
+    console_kwargs.update(console_extras)
+    
+  non_params_cols = 4
   rich_table = rich.table.Table(
       show_header=True,
       show_lines=True,
@@ -338,6 +337,8 @@ def _render_table(table: Table) -> str:
   )
 
   rich_table.add_column('path')
+  rich_table.add_column('module')
+  rich_table.add_column('inputs')
   rich_table.add_column('outputs')
 
   for col in table.collections:
@@ -351,20 +352,25 @@ def _render_table(table: Table) -> str:
 
       if collection in row.module_variables:
         col_repr += _as_yaml_str(
-            jax.tree_util.tree_map(_format_value,
-                                   row.module_variables[collection]))
-        col_repr += '\n\n'
+          _summary_tree_map(_ArrayRepresentation.render_array, row.module_variables[collection]))
+        if col_repr:
+          col_repr += '\n\n'
 
       col_repr += f'[bold]{_size_and_bytes_repr(*size_bytes)}[/bold]'
       collections_size_repr.append(col_repr)
 
+    no_show_methods = {'__call__', '<lambda>'}
+    path_repr = '/'.join(row.path)
+    method_repr = f' [dim]({row.method})[/dim]' if row.method not in no_show_methods else ''
     rich_table.add_row(
-        '/'.join(row.path) if row.path else table.module.__class__.__name__,
-        _as_yaml_str(jax.tree_util.tree_map(_format_value, row.outputs)),
+        path_repr,
+        row.module_type.__name__ + method_repr,
+        _as_yaml_str(_summary_tree_map(lambda x: x.render(), row.inputs)),
+        _as_yaml_str(_summary_tree_map(lambda x: x.render(), row.outputs)),
         *collections_size_repr)
 
   # add footer with totals
-  rich_table.columns[1].footer = rich.text.Text.from_markup(
+  rich_table.columns[non_params_cols - 1].footer = rich.text.Text.from_markup(
       'Total', justify='right')
 
   # get collection totals
@@ -378,8 +384,8 @@ def _render_table(table: Table) -> str:
 
   # add totals to footer
   for i, col in enumerate(table.collections):
-    rich_table.columns[2 +
-                       i].footer = _size_and_bytes_repr(*collection_total[col])
+    rich_table.columns[non_params_cols + i].footer = \
+      _size_and_bytes_repr(*collection_total[col])
 
   # add final totals to caption
   caption_totals = (0, 0)
@@ -392,8 +398,10 @@ def _render_table(table: Table) -> str:
   rich_table.caption_style = 'bold'
   rich_table.caption = f'\nTotal Parameters: {_size_and_bytes_repr(*caption_totals)}'
 
-  return '\n' + _get_rich_repr(rich_table) + '\n'
+  return '\n' + _get_rich_repr(rich_table, console_kwargs) + '\n'
 
+def _summary_tree_map(f, tree, *rest):
+  return jax.tree_util.tree_map(f, tree, *rest, is_leaf=lambda x: x is None)
 
 def _size_and_bytes_repr(size: int, num_bytes: int) -> str:
   if not size:
@@ -409,9 +417,9 @@ def _size_and_bytes(pytree: Any) -> Tuple[int, int]:
   return size, num_bytes
 
 
-def _get_rich_repr(obj):
+def _get_rich_repr(obj, console_kwargs):
   f = io.StringIO()
-  console = rich.console.Console(file=f, force_terminal=True)
+  console = rich.console.Console(file=f, **console_kwargs)
   console.print(obj)
   return f.getvalue()
 
@@ -432,13 +440,13 @@ def _as_yaml_str(value) -> str:
   return file.getvalue().replace('\n...', '').replace('\'', '').strip()
 
 
-def _format_value(value):
-  if hasattr(value, 'shape') and hasattr(value, 'dtype'):
-    shape_repr = ','.join(map(str, value.shape))
-    return f'[dim]{value.dtype}[/dim][{shape_repr}]'
+def _normalize_structure(obj):
+  if isinstance(obj, (tuple, list)):
+    return tuple(map(_normalize_structure, obj))
+  elif isinstance(obj, Mapping):
+    return {k: _normalize_structure(v) for k, v in obj.items()}
   else:
-    return str(value)
-
+    return obj
 
 def _bytes_repr(num_bytes):
   count, units = ((f'{num_bytes / 1e9 :,.1f}', 'GB') if num_bytes > 1e9 else


Test patch:
diff --git a/tests/linen/summary_test.py b/tests/linen/summary_test.py
--- a/tests/linen/summary_test.py
+++ b/tests/linen/summary_test.py
@@ -12,23 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import dataclasses
-from typing import List, Type
+from typing import List
 
 import jax
 import jax.numpy as jnp
-import numpy as np
 from absl.testing import absltest
-from jax import lax, random
-from jax.nn import initializers
+from jax import random
+import numpy as np
 
 from flax import linen as nn
 from flax.core.scope import Array
-from flax.linen.summary import _get_module_table
+from flax.linen import summary
 
 # Parse absl flags test_srcdir and test_tmpdir.
 jax.config.parse_flags_with_absl()
 
+CONSOLE_TEST_KWARGS = dict(force_terminal=False, no_color=True, width=10_000)
+
 def _get_shapes(pytree):
   return jax.tree_util.tree_map(lambda x: x.shape if hasattr(x, 'shape') else x, pytree)
 
@@ -96,9 +96,7 @@ def __call__(self, x: Array, training: bool) -> Array:
 
     return x, dict(a=x, b=x+1.0)
 
-
-
-class ModuleTest(absltest.TestCase):
+class SummaryTest(absltest.TestCase):
 
   def test_module_summary(self):
     """
@@ -111,55 +109,63 @@ def test_module_summary(self):
     x = jnp.ones((batch_size, 28, 28, 1))
     module = CNN(test_sow=False)
 
-    table = _get_module_table(
-      module, 
+    table = summary._get_module_table(module, depth=None, show_repeated=True)( 
       {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)},
-      method=None, mutable=True, depth=None,
-      exclude_methods=set(),
-    )( 
-      x, training=True
+      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(lambda x: x.value(), row.inputs)
+      row.outputs = jax.tree_util.tree_map(lambda x: x.value(), row.outputs)
 
-    # 11 rows = 1 Inputs + 4 ConvBlock_0 + 4 ConvBlock_1 + 1 Dense_0 + 1 Module output
-    self.assertEqual(len(table), 11)
+    # 10 rows = 1 CNN + 4 ConvBlock_0 + 4 ConvBlock_1 + 1 Dense_0
+    self.assertEqual(len(table), 10)
 
     # check paths
-    self.assertEqual(table[0].path, ("Inputs",))
+    self.assertEqual(table[0].path, ())
 
-    self.assertEqual(table[1].path, ("block1", "bn"))
+    self.assertEqual(table[1].path, ("block1",))
     self.assertEqual(table[2].path, ("block1", "conv"))
-    self.assertEqual(table[3].path, ("block1", "dropout"))
-    self.assertEqual(table[4].path, ("block1",))
+    self.assertEqual(table[3].path, ("block1", "bn"))
+    self.assertEqual(table[4].path, ("block1", "dropout"))
 
-    self.assertEqual(table[5].path, ("block2", "bn"))
+    self.assertEqual(table[5].path, ("block2",))
     self.assertEqual(table[6].path, ("block2", "conv"))
-    self.assertEqual(table[7].path, ("block2", "dropout"))
-    self.assertEqual(table[8].path, ("block2",))
+    self.assertEqual(table[7].path, ("block2", "bn"))
+    self.assertEqual(table[8].path, ("block2", "dropout"))
 
     self.assertEqual(table[9].path, ("dense",))
-    self.assertEqual(table[10].path, ())
 
     # check outputs shapes
     self.assertEqual(
-      (table[0].outputs[0].shape, table[0].outputs[1]),
+      (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))
-    self.assertEqual(
-      _get_shapes(table[10].outputs),
-      ((batch_size, 10), dict(a=(batch_size, 10), b=(batch_size, 10))),
-    )
 
     # check no summary is performed
     for row in table:
@@ -178,45 +184,51 @@ def test_module_summary_with_depth(self):
     x = jnp.ones((batch_size, 28, 28, 1))
     module = CNN(test_sow=False)
 
-    table = _get_module_table(
-      module, 
+    table = summary._get_module_table(module, depth=1, show_repeated=True)(
       {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)},
-      method=None, mutable=True, depth=1,
-      exclude_methods=set(),
-    )(
-      x, training=True
+      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(lambda x: x.value(), row.inputs)
+      row.outputs = jax.tree_util.tree_map(lambda x: x.value(), row.outputs)
 
-    # 5 rows = 1 Inputs + 1 ConvBlock_0 + 1 ConvBlock_1 + 1 Dense_0 + 1 Module output
-    self.assertEqual(len(table), 5)
+    # 4 rows = 1 CNN + 1 ConvBlock_0 + 1 ConvBlock_1 + 1 Dense_0
+    self.assertEqual(len(table), 4)
 
     # check paths
-    self.assertEqual(table[0].path, ("Inputs",))
+    self.assertEqual(table[0].path, ())
+
     self.assertEqual(table[1].path, ("block1",))
     self.assertEqual(table[2].path, ("block2",))
     self.assertEqual(table[3].path, ("dense",))
-    self.assertEqual(table[4].path, ())
 
     # check outputs shapes
     self.assertEqual(
-      (table[0].outputs[0].shape, table[0].outputs[1]),
+      (table[0].inputs[0].shape, table[0].inputs[1]),
       (x.shape, dict(training=True)),
     )
-    self.assertEqual(table[1].outputs.shape, (batch_size, 28, 28, 32))
-    self.assertEqual(table[2].outputs.shape, (batch_size, 28, 28, 64))
-    self.assertEqual(table[3].outputs.shape, (batch_size, 10))
     self.assertEqual(
-      _get_shapes(table[4].outputs),
+      _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 Dense_0 and Module output are not summarized
+    # 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)
-    self.assertEqual(table[4].module_variables, table[4].counted_variables)
 
   
   def test_tabulate(self):
@@ -233,6 +245,7 @@ def test_tabulate(self):
         {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)}, 
         x, 
         training=True,
+        console_kwargs=CONSOLE_TEST_KWARGS,
     )
 
     # NOTE: its tricky to validate the content of lines
@@ -246,9 +259,11 @@ def test_tabulate(self):
 
     # 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("batch_stats", lines[3])
+    self.assertIn("batch_stats", lines[3]) 
 
     # collection counts
     self.assertIn("Total", lines[-6])
@@ -274,9 +289,11 @@ def test_tabulate_with_sow(self):
       {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)}, 
       x, 
       training=True,
+      console_kwargs=CONSOLE_TEST_KWARGS,
     )
 
-    self.assertNotIn("INTERM", module_repr)
+    self.assertIn("intermediates", module_repr)
+    self.assertIn("INTERM", module_repr)
   
   def test_tabulate_with_method(self):
 
@@ -290,9 +307,11 @@ def test_tabulate_with_method(self):
       x, 
       training=True,
       method=CNN.cnn_method,
+      console_kwargs=CONSOLE_TEST_KWARGS,
     )
 
-    self.assertNotIn("INTERM", module_repr)
+    self.assertIn("(block_method)", module_repr)
+    self.assertIn("(cnn_method)", module_repr)
 
   def test_tabulate_function(self):
     """
@@ -307,14 +326,12 @@ def test_tabulate_function(self):
     module_repr = nn.tabulate(
       module,
       {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)},
+      console_kwargs=CONSOLE_TEST_KWARGS,
     )(
       x,
       training=True,
     )
 
-    # NOTE: its tricky to validate the content of lines
-    # because it seems to be shell-dependent, so we will
-    # just check lines that wont change between environments
     lines = module_repr.split("\n")
 
     # check title
@@ -323,6 +340,8 @@ def test_tabulate_function(self):
 
     # 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("batch_stats", lines[3])
@@ -337,4 +356,158 @@ def test_tabulate_function(self):
     # total counts
     self.assertIn("Total Parameters", lines[-3])
     self.assertIn("19,850", lines[-3])
-    self.assertIn("79.4 KB", lines[-3])
\ No newline at end of file
+    self.assertIn("79.4 KB", lines[-3])
+
+
+  def test_lifted_transform(self):
+    class LSTM(nn.Module):
+      batch_size: int
+      out_feat: int
+
+      @nn.compact
+      def __call__(self, x):
+          carry = nn.LSTMCell.initialize_carry(
+              random.PRNGKey(0), (self.batch_size,), self.out_feat
+          )
+          Cell = nn.scan(
+              nn.LSTMCell,
+              variable_broadcast="params",
+              split_rngs={"params": False},
+              in_axes=1,
+              out_axes=1,
+          )
+          return Cell(name="ScanLSTM")(carry, x)
+
+
+    lstm = LSTM(batch_size=32, out_feat=128)
+
+    with jax.check_tracer_leaks(True):
+      module_repr = lstm.tabulate(
+        random.PRNGKey(0),
+        x=jnp.ones((32, 128, 64)), 
+        console_kwargs=CONSOLE_TEST_KWARGS)
+    
+    lines = module_repr.splitlines()
+
+    self.assertIn("LSTM", lines[5])
+    self.assertIn("ScanLSTM", lines[9])
+    self.assertIn("LSTMCell", lines[9])
+    self.assertIn("ScanLSTM/ii", lines[13])
+    self.assertIn("Dense", lines[13])
+  
+  def test_lifted_transform_no_rename(self):
+    class LSTM(nn.Module):
+      batch_size: int
+      out_feat: int
+
+      @nn.compact
+      def __call__(self, x):
+          carry = nn.LSTMCell.initialize_carry(
+              random.PRNGKey(0), (self.batch_size,), self.out_feat
+          )
+          Cell = nn.scan(
+              nn.LSTMCell,
+              variable_broadcast="params",
+              split_rngs={"params": False},
+              in_axes=1,
+              out_axes=1,
+          )
+          return Cell()(carry, x)
+
+
+    lstm = LSTM(batch_size=32, out_feat=128)
+
+    with jax.check_tracer_leaks(True):
+      module_repr = lstm.tabulate(
+        random.PRNGKey(0),
+        x=jnp.ones((32, 128, 64)), 
+        console_kwargs=CONSOLE_TEST_KWARGS)
+
+    lines = module_repr.splitlines()
+
+    self.assertIn("LSTM", lines[5])
+    self.assertIn("ScanLSTMCell_0", lines[9])
+    self.assertIn("LSTMCell", lines[9])
+    self.assertIn("ScanLSTMCell_0/ii", lines[13])
+    self.assertIn("Dense", lines[13])
+
+  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.PRNGKey(0), 
+      x=x,
+      show_repeated=True,
+      console_kwargs=CONSOLE_TEST_KWARGS)
+    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])
+
+  def test_empty_input(self):
+    class EmptyInput(nn.Module):
+      @nn.compact
+      def __call__(self):
+        return 1
+
+    module = EmptyInput()
+    module_repr = module.tabulate({}, console_kwargs=CONSOLE_TEST_KWARGS)
+    lines = module_repr.splitlines()
+
+    self.assertRegex(lines[5], r'|\s*|\s*EmptyInput\s*|\s*|\s*1\s*|')
+  
+  def test_numpy_scalar(self):
+    class Submodule(nn.Module):
+      def __call__(self, x):
+        return x + 1
+
+    class EmptyInput(nn.Module):
+      @nn.compact
+      def __call__(self):
+        return Submodule()(x=np.pi)
+
+    module = EmptyInput()
+    module_repr = module.tabulate({}, console_kwargs=CONSOLE_TEST_KWARGS)
+    lines = module_repr.splitlines()
+
+    self.assertIn('4.141592', lines[5])
+    self.assertIn('x: 3.141592', lines[7])
+    self.assertIn('4.141592', lines[7])
\ No newline at end of file