google / google/flax

Improve Error Message: Assigning a list of Modules to `self` incorrectly

Open
#524 15 comments 3 reactions 2 assignees Claimed by @marcvanzee View on GitHub
Priority: P1 - soon
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

In `setup()`, submodules are expected to "know their name" when they are assigned to `self`, which is implemented by overriding `__setattr__`. This can cause problems when appending modules to a list. Consider the following code.

```
class Test(nn.Module):
def setup(self):
self.layers = []
for i in range(3):
self.layers.append(nn.Dense(5))

def __call__(self, x):
for layer in self.layers:
x = layer(x) # <--- ValueError: Can't call methods on orphaned modules
return x
```

The code is failing because `self.layers.append()` does not assign anything to `self`, so the overwritten `__setattr__` isn't triggered. As a result, the `Dense` modules are not properly initialized and they are recognized as "orphaned" (no parent scope). The canonical way to fix this code is:
```
def setup(self):
self.layers = [nn.Dense(5) for _ in range(3)]
```

Requiring the user to understand this distinction causes mental overload, especially since the error message is cryptic.

Some suggestions for improvements:

1. Improve the error message by explaining how modules should be assigned to `self` in `setup`.

2. Create a new abstraction `ModuleList`, and only allow lists of Modules here (similar to PyTorch).

3. Convert the list to a some data structure that tracks mutation properly.

Option (1) is a quick fix, but doesn't seem like a great long-term solution, since this error can appear in other cases as well, so we want to make sure the error message also covers the other cases.

Option (2) doesn't require much mental overload, but the user should remember that this abstraction exists.

Option (3) would also work, but I fear that this magic may lead to new special cases in the future that we should address.

I personally feel option (2) is most consistent and robust in the long-term. Since it doesn't have to be used in a particular way, it seems also easiest to understand for users.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.