Repository: google/flax

Issue:
Error when calling module tabulate involving WeightNorm
### System information
- WSL on Windows 11 Pro
- Flax, jax, jaxlib versions:
```
Name: flax
Version: 0.8.1
Summary: Flax: A neural network library for JAX designed for flexibility
Home-page:
Author:
Author-email: Flax team <flax-dev@google.com>
License:
Location: /home/admin/.local/lib/python3.10/site-packages
Requires: jax, msgpack, numpy, optax, orbax-checkpoint, PyYAML, rich, tensorstore, typing-extensions
Required-by: clu
---
Name: jax
Version: 0.4.24
Summary: Differentiate, compile, and transform Numpy code.
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/admin/.local/lib/python3.10/site-packages
Requires: ml-dtypes, numpy, opt-einsum, scipy
Required-by: chex, clu, flax, jaxloudnorm, optax, orbax-checkpoint
---
Name: jaxlib
Version: 0.4.24+cuda12.cudnn89
Summary: XLA library for JAX
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/admin/.local/lib/python3.10/site-packages
Requires: ml-dtypes, numpy, scipy
Required-by: chex, clu, optax, orbax-checkpoint
```
- Python version: 3.10
- GPU/TPU model and memory: Nvidia RTX 2080 (8 GB)
- CUDA version (if applicable): 12.3

### Discussed in https://github.com/google/flax/discussions/3700

<div type='discussions-op-text'>

<sup>Originally posted by **DBraun** February 19, 2024</sup>
I'm unable to call `model.tabulate` with `compute_flops=True` or  `compute_vjp_flops=True` if the module involves `flax.linen.WeightNorm`.

Code:
```python
import jax
import jax.numpy as jnp
from flax import linen as nn

model = nn.WeightNorm(nn.Conv(features=64, kernel_size=3, strides=2))
# model = nn.Conv(features=64, kernel_size=3, strides=2)

x = jnp.ones((8, 44100, 1))

key = jax.random.PRNGKey(0)
params = model.init(key, x)

print(model.tabulate(key, x, depth=4,
                     console_kwargs={'width': 180},
                     column_kwargs={'width': 180},
                     compute_flops=True,
                     compute_vjp_flops=True,
                     ))

y = model.apply(params, x)
print(y.shape)
```

Output:
```
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/mnt/c/Users/admin/AppData/Roaming/JetBrains/PyCharm2023.2/scratches/scratch1.py", line 13, in <module>
    print(model.tabulate(key, x, depth=4,
  File "/home/admin/.local/lib/python3.10/site-packages/flax/linen/module.py", line 2692, in tabulate
    return tabulate_fn(*args, **kwargs)
  File "/home/admin/.local/lib/python3.10/site-packages/flax/linen/summary.py", line 315, in _tabulate_fn
    table = table_fn(rngs, *fn_args, **fn_kwargs, **kwargs)
  File "/home/admin/.local/lib/python3.10/site-packages/flax/linen/summary.py", line 488, in _get_table_fn
    *_get_call_flops(c, compute_flops, compute_vjp_flops),
  File "/home/admin/.local/lib/python3.10/site-packages/flax/linen/summary.py", line 398, in _get_call_flops
    variables = jax.eval_shape(init, rngs, dynamic_leaves)
  File "/home/admin/.local/lib/python3.10/site-packages/flax/linen/summary.py", line 390, in init
    return c.module.init(
AttributeError: "Conv" object has no attribute "layer_forward". If "layer_forward" is defined in '.setup()', remember these fields are only accessible from inside 'init' or 'apply'.
```

Maybe the answer involves specifying `method=...` in the call to `tabulate` but I haven't figured it out yet. What's the solution?

For context, the model that I want is based off the following PyTorch code:
```python
import torch.nn as nn
from torch.nn.utils import weight_norm
model = weight_norm(nn.Conv1d(1, 64, 3))
```</div>


Discussion:
@chiamp 
Seems like `init` is called on `Conv` (`c.module`) using the method `layer_forward` ([`c.method`](https://github.com/google/flax/blob/main/flax/linen/summary.py#L395)), but I'm not sure why there is a `_CallInfo` object with these attributes since `layer_forward` is a `WeightNorm` method, not a `Conv` method. Any thoughts @cgarciae?
When computing flops each Module is run under `eval_shape`, there is a weird intereaction between `WeightNorm` and the tabulate context here as @chiamp points out.
Thanks for looking into it. One more thing, hopefully the solution will also fix the similar situation for `ConvTranspose` :
```python
model = nn.WeightNorm(nn.ConvTranspose(features=64, kernel_size=(3,), strides=(2,), transpose_kernel=True))
```

Resolving patch:
diff --git a/flax/linen/summary.py b/flax/linen/summary.py
--- a/flax/linen/summary.py
+++ b/flax/linen/summary.py
@@ -460,7 +460,7 @@ def _get_variables():
       call_depth = len(c.path)
       inputs = _process_inputs(c.args, c.kwargs)
 
-      if c.path in visited_paths:
+      if c.path in visited_paths or not hasattr(c.module, c.method):
         if not show_repeated:
           continue
         module_vars = {}


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
@@ -756,6 +756,57 @@ class InputEnum(str, enum.Enum):
     self.assertIn('inputs:', lines[5])
     self.assertIn('x: \x1b[2mfloat32\x1b[0m[1,1]', lines[6])
 
+  def test_tabulate_norm_wrapper(self):
+    class SubModel(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        x = nn.SpectralNorm(nn.Dense(5))(x, update_stats=False)
+        x = nn.Dense(6)(x)
+        x = nn.WeightNorm(nn.Dense(7))(x)
+        return x
+
+    class Model(nn.Module):
+      @nn.compact
+      def __call__(self, x):
+        x = nn.WeightNorm(nn.Dense(3))(x)
+        x = nn.Dense(4)(x)
+        x = SubModel()(x)
+        x = nn.Dense(8)(x)
+        x = nn.SpectralNorm(nn.Dense(9))(x, update_stats=False)
+        return x
+
+    x = jnp.ones((1, 2))
+    key = jax.random.key(0)
+    model = Model()
+
+    lines = model.tabulate(
+      key,
+      x,
+      console_kwargs=CONSOLE_TEST_KWARGS,
+      compute_flops=True,
+      compute_vjp_flops=True,
+    ).splitlines()
+
+    self.assertIn('Model', lines[5])
+    self.assertIn('WeightNorm_0', lines[7])
+    self.assertIn('Dense_0/kernel/scale', lines[7])
+    self.assertIn('Dense_0', lines[11])
+    self.assertIn('Dense_1', lines[16])
+    self.assertIn('SubModel_0', lines[21])
+    self.assertIn('SubModel_0/SpectralNorm_0', lines[23])
+    self.assertIn('Dense_0/kernel/sigma', lines[23])
+    self.assertIn('Dense_0/kernel/u', lines[24])
+    self.assertIn('SubModel_0/Dense_0', lines[28])
+    self.assertIn('SubModel_0/Dense_1', lines[33])
+    self.assertIn('SubModel_0/WeightNorm_0', lines[38])
+    self.assertIn('Dense_2/kernel/scale', lines[38])
+    self.assertIn('SubModel_0/Dense_2', lines[42])
+    self.assertIn('Dense_2', lines[47])
+    self.assertIn('SpectralNorm_0', lines[52])
+    self.assertIn('Dense_3/kernel/sigma', lines[52])
+    self.assertIn('Dense_3/kernel/u', lines[53])
+    self.assertIn('Dense_3', lines[57])
+
 
 if __name__ == '__main__':
   absltest.main()