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

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.

[cgarciae]

After reading all the thread I got an alternative proposal:

• No ModuleList/Dict abstraction, users just use regular python lists/dicts.
• The list -> tuple and dict -> FrozenDict conversion happens after setup finishes, this way the user can mutate the list/dicts inside setup.
• __setattr__ only keeps track of the set of fields that need to be frozen instead of doing on the spot. Alternatively, just iterate over all fields and check which need to be frozen after setup.

[avital]

I am sympathetic to option (2), but here are a few questions:

1. Could/should we raise an error if you assign a “vanilla” list of Modules?
2. What happens if you want to store a list that has both Modules and other things (e.g. the relu function, which is a common need for some Sequential patterns)

Also curious to hear from @levskaya who wrote the original “module list logic” in Modules.