Improve Error Message: Assigning a list of Modules to `self` incorrectly
- 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
Assessment
This issue has not been assessed yet.