Repository: google/flax

Issue:
Tabulate having issues with T5X
This was reported by Boris offline, will try to get a repro.


Discussion:

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -65,17 +65,6 @@
 
 # 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
   lines = x.split('\n')
@@ -836,8 +825,7 @@ def _call_wrapped_method(self, fun, args, kwargs):
         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)
+        _args, _kwargs, _y = flax.linen.summary._represent_tree((args, kwargs, y))
         _context.call_info_stack[-1].calls.append(
           _CallInfo(call_index, scope_path, type(self), fun.__name__, _args, _kwargs, _y))
       return y
diff --git a/flax/linen/summary.py b/flax/linen/summary.py
--- a/flax/linen/summary.py
+++ b/flax/linen/summary.py
@@ -19,6 +19,7 @@
 from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union
 
 import flax.linen.module as module_lib
+from flax.core import meta
 from flax.core.scope import CollectionFilter, FrozenVariableDict, MutableVariableDict
 import jax
 import jax.numpy as jnp
@@ -26,6 +27,7 @@
 import rich.table
 import rich.text
 import yaml
+import numpy as np
 
 PRNGKey = Any  # pylint: disable=invalid-name
 RNGSequences = Dict[str, PRNGKey]
@@ -38,25 +40,34 @@ class _ValueRepresentation(ABC):
   def render(self) -> str:
     ...
 
-  @abstractmethod
-  def value(self) -> Any:
-    ...
-
 @dataclasses.dataclass
 class _ArrayRepresentation(_ValueRepresentation):
   shape: Tuple[int, ...]
   dtype: Any
 
+  @classmethod
+  def from_array(cls, x: Array) -> '_ArrayRepresentation':
+    return cls(jnp.shape(x), jnp.result_type(x))
+
   @classmethod
   def render_array(cls, x) -> str:
-    return cls(jnp.shape(x), jnp.result_type(x)).render()
+    return cls.from_array(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 _PartitionedArrayRepresentation(_ValueRepresentation):
+  array_representation: _ArrayRepresentation
+  names: meta.LogicalNames
+
+  @classmethod
+  def from_partitioned(cls, partitioned: meta.Partitioned) -> '_PartitionedArrayRepresentation':
+    return cls(_ArrayRepresentation.from_array(partitioned.value), partitioned.names)
+
+  def render(self):
+    return self.array_representation.render() + f' [dim]P[/dim]{self.names}'
 
 @dataclasses.dataclass
 class _ObjectRepresentation(_ValueRepresentation):
@@ -65,9 +76,6 @@ class _ObjectRepresentation(_ValueRepresentation):
   def render(self):
     return repr(self.obj)
 
-  def value(self):
-    return self.obj
-
 @dataclasses.dataclass
 class Row:
   """Contains the information about a single row in the summary table.
@@ -92,8 +100,10 @@ class Row:
   counted_variables: Dict[str, Dict[str, Any]]
 
   def __post_init__(self):
-    self.inputs = _normalize_structure(self.inputs)
-    self.outputs = _normalize_structure(self.outputs)
+    self.inputs = self.inputs
+    self.outputs = self.outputs
+    self.module_variables = self.module_variables
+    self.counted_variables = self.counted_variables
 
   def size_and_bytes(self, collections: Iterable[str]) -> Dict[str, Tuple[int, int]]:
     return {
@@ -355,8 +365,10 @@ def _render_table(table: Table, console_extras: Optional[Mapping[str, Any]]) ->
       col_repr = ''
 
       if collection in row.module_variables:
+        module_variables = _represent_tree(row.module_variables[collection])
+        module_variables = _normalize_structure(module_variables)
         col_repr += _as_yaml_str(
-          _summary_tree_map(_ArrayRepresentation.render_array, row.module_variables[collection]))
+          _summary_tree_map(_maybe_render, module_variables))
         if col_repr:
           col_repr += '\n\n'
 
@@ -369,8 +381,8 @@ def _render_table(table: Table, console_extras: Optional[Mapping[str, Any]]) ->
     rich_table.add_row(
         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)),
+        _as_yaml_str(_summary_tree_map(_maybe_render, _normalize_structure(row.inputs))),
+        _as_yaml_str(_summary_tree_map(_maybe_render, _normalize_structure(row.outputs))),
         *collections_size_repr)
 
   # add footer with totals
@@ -416,8 +428,8 @@ def _size_and_bytes_repr(size: int, num_bytes: int) -> str:
 
 def _size_and_bytes(pytree: Any) -> Tuple[int, int]:
   leaves = jax.tree_util.tree_leaves(pytree)
-  size = sum(x.size for x in leaves)
-  num_bytes = sum(x.size * x.dtype.itemsize for x in leaves)
+  size = sum(x.size for x in leaves if hasattr(x, 'size'))
+  num_bytes = sum(x.size * x.dtype.itemsize for x in leaves if hasattr(x, 'size'))
   return size, num_bytes
 
 
@@ -445,10 +457,14 @@ def _as_yaml_str(value) -> str:
 
 
 def _normalize_structure(obj):
+  if isinstance(obj, _ValueRepresentation):
+    return 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()}
+  elif dataclasses.is_dataclass(obj):
+    return {f.name: _normalize_structure(getattr(obj, f.name)) for f in dataclasses.fields(obj)}
   else:
     return obj
 
@@ -459,3 +475,25 @@ def _bytes_repr(num_bytes):
                   (f'{num_bytes:,}', 'B'))
 
   return f'{count} {units}'
+
+
+def _get_value_representation(x: Any) -> _ValueRepresentation:
+  if isinstance(x, (int, float, bool, type(None))) or (
+    isinstance(x, np.ndarray) and np.isscalar(x)):
+    return _ObjectRepresentation(x)
+  elif isinstance(x, meta.Partitioned):
+    return _PartitionedArrayRepresentation.from_partitioned(x)
+  try:
+    return _ArrayRepresentation.from_array(x)
+  except:
+    return _ObjectRepresentation(x)
+
+def _represent_tree(x):
+  """Returns a tree with the same structure as `x` but with each leaf replaced
+  by a `_ValueRepresentation` object."""
+  return jax.tree_util.tree_map(
+    _get_value_representation, x,
+    is_leaf=lambda x: x is None or isinstance(x, meta.Partitioned))
+
+def _maybe_render(x):
+  return x.render() if hasattr(x, 'render') else repr(x)
\ No newline at end of file


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
@@ -23,6 +23,7 @@
 from flax import linen as nn
 from flax.core.scope import Array
 from flax.linen import summary
+from flax import struct
 
 # Parse absl flags test_srcdir and test_tmpdir.
 jax.config.parse_flags_with_absl()
@@ -32,6 +33,11 @@
 def _get_shapes(pytree):
   return jax.tree_util.tree_map(lambda x: x.shape if hasattr(x, 'shape') else x, pytree)
 
+def _get_obj_repr_value(x):
+  if isinstance(x, summary._ObjectRepresentation):
+    return x.obj
+  return x
+
 class ConvBlock(nn.Module):
   features: int
   kernel_size: List[int]
@@ -115,8 +121,8 @@ def test_module_summary(self):
     )
     # 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)
+      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.assertEqual(len(table), 10)
@@ -189,9 +195,10 @@ def test_module_summary_with_depth(self):
       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)
+      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.assertEqual(len(table), 4)
@@ -512,6 +519,50 @@ def __call__(self):
     self.assertIn('x: 3.141592', lines[7])
     self.assertIn('4.141592', lines[7])
 
+  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.PRNGKey(0), jnp.empty((1, 28, 28, 1)),
+                            console_kwargs=CONSOLE_TEST_KWARGS).splitlines()
+    self.assertIn('P(None,)', lines[7])
+    self.assertIn('P(None, data)', lines[8])
+
+  def test_non_array_variables(self):
+
+    class Metadata(struct.PyTreeNode):
+      names: tuple = struct.field(pytree_node=False)
+
+    class Foo(nn.Module):
+      @nn.compact
+      def __call__(self):
+        self.sow('foo', 'bar', Metadata(('baz', 'qux')))
+
+    module = Foo()
+    lines = module.tabulate({},
+                            console_kwargs=CONSOLE_TEST_KWARGS).splitlines()
+    self.assertIn('names', lines[6])
+    self.assertIn('baz', lines[7])
+    self.assertIn('qux', lines[8])
+
 
 if __name__ == '__main__':
   absltest.main()