Error when calling `Module.tabulate` on normalization wrappers like `WeightNorm` and `SpectralNorm`
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
Follow-up from #3735. Partial fix in #3772.
Minimum repro:
```
import jax, jax.numpy as jnp
from flax import linen as nn
model = nn.WeightNorm(nn.Dense(3))
x = jnp.ones((1, 2))
key = jax.random.key(0)
print(model.tabulate(key,
x,
compute_flops=True,
compute_vjp_flops=True,
))
```
Error message:
```
Traceback (most recent call last):
File "/Users/marcuschiam/Desktop/asdf.py", line 20, in
print(model.tabulate(key,
File "/Users/marcuschiam/flax/flax/linen/module.py", line 2843, in tabulate
return tabulate_fn(*args, **kwargs)
File "/Users/marcuschiam/flax/flax/linen/summary.py", line 315, in _tabulate_fn
table = table_fn(rngs, *fn_args, **fn_kwargs, **kwargs)
File "/Users/marcuschiam/flax/flax/linen/summary.py", line 490, in _get_table_fn
*_get_call_flops(c, compute_flops, compute_vjp_flops),
File "/Users/marcuschiam/flax/flax/linen/summary.py", line 400, in _get_call_flops
variables = jax.eval_shape(init, rngs, dynamic_leaves)
File "/Users/marcuschiam/flax/flax/linen/summary.py", line 392, in init
return c.module.init(
AttributeError: "Dense" object has no attribute "layer_forward". If "layer_forward" is defined in '.setup()', remember these fields are only accessible from inside 'init' or 'apply'.
```
When calling `tabulate()`, [`calls`](https://github.com/google/flax/blob/main/flax/linen/summary.py#L451) contains [`_CallInfo`](https://github.com/google/flax/blob/main/flax/linen/module.py#L155) objects where the [`method`](https://github.com/google/flax/blob/main/flax/linen/module.py#L161) does not exist in the [`module`](https://github.com/google/flax/blob/main/flax/linen/module.py#L158), causing the [`AttributeError`](https://github.com/google/flax/blob/main/flax/linen/summary.py#L395) seen in #3735.
This is illustrated by inspecting the `_CallInfo` objects in the debugger:
```
for c in calls:
print(f'Index: {c.index}\tModule: {type(c.module)}\tPath: {c.path}\tMethod: {c.method}')
```
```
Index: 0 Module: Path: () Method: __call__
Index: 1 Module: Path: () Method:
Index: 2 Module: Path: ('layer_instance',) Method: layer_forward
Index: 3 Module: Path: ('layer_instance',) Method: __call__
Index: 4 Module: Path: () Method: _l2_normalize
Index: 5 Module: Path: () Method: _l2_normalize
Index: 6 Module: Path: ('layer_instance',) Method: layer_forward
Index: 7 Module: Path: ('layer_instance',) Method: __call__
```
For the `_CallInfo` object with index 2, `Dense` does not have a `layer_forward` method/attribute and an error is thrown.
For the `_CallInfo` object with index 1, the `` function is skipped because the path `()` has [already been visited](https://github.com/google/flax/blob/main/flax/linen/summary.py#L463) in index 0.
Contributor guide
Assessment
This issue has not been assessed yet.