Repository: google/flax

Issue:
Can't tabulate a linen module that wants to be pmapped
I have a custom `flax.linen.Module` that uses collective operations (like `jax.lax.psum`) and therefore has to be under `pmap` to be used. I noticed that `tabulate()` doesn't work anymore if I apply `pmap` to the module's `init` function.

Example:

```python
import jax
import jax.numpy as jnp
import flax.linen as nn

key = jax.random.PRNGKey(0)

class MyModule(nn.Module):
    @nn.compact
    def __call__(self, inputs):
        return nn.Dense(100)(inputs)
    
# tabulate without pmap: all good
model = MyModule()
print(model.tabulate(key, jnp.ones((100, 100))))

# apply pmap
model.init = jax.pmap(model.init, axis_name="i")
replicated_inputs = (jnp.stack([key] * 4), jnp.ones((4, 100, 100)))

# can still use eval_shape
print(jax.eval_shape(model.init, *replicated_inputs))

# tabulate doesn't work anymore
print(model.tabulate(*replicated_inputs))
```

Traceback:

```python
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
File /anaconda/envs/py311/lib/python3.11/site-packages/jax/_src/api.py:1278, in _mapped_axis_size.<locals>._get_axis_size(name, shape, axis)
   1277 try:
-> 1278   return shape[axis]
   1279 except (IndexError, TypeError) as e:

IndexError: tuple index out of range

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
Cell In[43], line 24
     21 print(jax.eval_shape(model.init, *replicated_inputs))
     23 # nope
---> 24 print(model.tabulate(*replicated_inputs))

File /anaconda/envs/py311/lib/python3.11/site-packages/flax/linen/module.py:2208, in Module.tabulate(self, rngs, depth, show_repeated, mutable, console_kwargs, table_kwargs, column_kwargs, *args, **kwargs)
   2196 from flax.linen import summary
   2198 tabulate_fn = summary.tabulate(
   2199     self,
   2200     rngs,
   (...)
   2206     column_kwargs=column_kwargs,
   2207 )
-> 2208 return tabulate_fn(*args, **kwargs)

File /anaconda/envs/py311/lib/python3.11/site-packages/flax/linen/summary.py:259, in tabulate.<locals>._tabulate_fn(*fn_args, **fn_kwargs)
    255 def _tabulate_fn(*fn_args, **fn_kwargs):
    256   table_fn = _get_module_table(
    257       module, depth=depth, show_repeated=show_repeated
    258   )
--> 259   table = table_fn(rngs, *fn_args, mutable=mutable, **fn_kwargs, **kwargs)
    260   return _render_table(table, console_kwargs, table_kwargs, column_kwargs)

File /anaconda/envs/py311/lib/python3.11/site-packages/flax/linen/summary.py:279, in _get_module_table.<locals>._get_table_fn(*args, **kwargs)
    276 def _get_variables():
    277   return module.init(*args, **kwargs)
--> 279 variables = jax.eval_shape(_get_variables)
    280 calls = module_lib._context.call_info_stack[-1].calls
    281 calls.sort(key=lambda c: c.index)

    [... skipping hidden 8 frame]

File /anaconda/envs/py311/lib/python3.11/site-packages/flax/linen/summary.py:277, in _get_module_table.<locals>._get_table_fn.<locals>._get_variables()
    276 def _get_variables():
--> 277   return module.init(*args, **kwargs)

    [... skipping hidden 7 frame]

File /anaconda/envs/py311/lib/python3.11/site-packages/jax/_src/api.py:1282, in _mapped_axis_size.<locals>._get_axis_size(name, shape, axis)
   1280 min_rank = axis + 1 if axis >= 0 else -axis
   1281 # TODO(mattjj): better error message here
-> 1282 raise ValueError(
   1283     f"{name} was requested to map its argument along axis {axis}, "
   1284     f"which implies that its rank should be at least {min_rank}, "
   1285     f"but is only {len(shape)} (its shape is {shape})") from e

ValueError: pmap was requested to map its argument along axis 0, which implies that its rank should be at least 1, but is only 0 (its shape is ())
```


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
@@ -2357,7 +2357,7 @@ def tabulate(
       *args,
       depth: Optional[int] = None,
       show_repeated: bool = False,
-      mutable: CollectionFilter = True,
+      mutable: CollectionFilter = DenyList('intermediates'),
       console_kwargs: Optional[Mapping[str, Any]] = None,
       table_kwargs: Mapping[str, Any] = MappingProxyType({}),
       column_kwargs: Mapping[str, Any] = MappingProxyType({}),
diff --git a/flax/linen/summary.py b/flax/linen/summary.py
--- a/flax/linen/summary.py
+++ b/flax/linen/summary.py
@@ -22,7 +22,7 @@
 
 import flax.linen.module as module_lib
 from flax.core import meta
-from flax.core.scope import CollectionFilter, FrozenVariableDict, MutableVariableDict
+from flax.core.scope import CollectionFilter, DenyList, FrozenVariableDict, MutableVariableDict
 import jax
 import jax.numpy as jnp
 import rich.console
@@ -156,7 +156,7 @@ def tabulate(
     rngs: Union[PRNGKey, RNGSequences],
     depth: Optional[int] = None,
     show_repeated: bool = False,
-    mutable: CollectionFilter = True,
+    mutable: CollectionFilter = DenyList('intermediates'),
     console_kwargs: Optional[Mapping[str, Any]] = None,
     table_kwargs: Mapping[str, Any] = MappingProxyType({}),
     column_kwargs: Mapping[str, Any] = MappingProxyType({}),
@@ -251,12 +251,16 @@ def __call__(self, x):
     (`method`) and returns a string with a tabular representation of the
     Modules.
   """
+  # add non-default arguments to kwargs, this prevents some issue we overloading init
+  # see: https://github.com/google/flax/issues/3299
+  if mutable != DenyList('intermediates'):
+    kwargs['mutable'] = mutable
 
   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)
+    table = table_fn(rngs, *fn_args, **fn_kwargs, **kwargs)
     return _render_table(table, console_kwargs, table_kwargs, column_kwargs)
 
   return _tabulate_fn


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
@@ -329,6 +329,7 @@ def test_tabulate_with_sow(self):
         {"dropout": random.key(0), "params": random.key(1)},
         x,
         training=True,
+        mutable=True,
         console_kwargs=CONSOLE_TEST_KWARGS,
     )
 
@@ -613,6 +614,27 @@ def __call__(self, x):
     lines = rep.splitlines()
     self.assertIn("Total Parameters: 50", lines[-2])
 
+  def test_monkey_patching_init(self):
+    class MyModule(nn.Module):
+
+      @nn.compact
+      def __call__(self, inputs):
+        return nn.Dense(3)(inputs)
+
+    key = jax.random.PRNGKey(0)
+
+    # tabulate without pmap
+    model = MyModule()
+    model.tabulate(key, jnp.ones((3, 3)))
+
+    # apply pmap
+    model.init = jax.pmap(model.init, axis_name="i")
+    rngs = jax.random.split(key, 4)
+    x = jnp.ones((4, 3, 3))
+
+    # test tabulate still works
+    model.tabulate(rngs, x)
+
 
 if __name__ == "__main__":
   absltest.main()