Linen: error with self.name in setup function

[Daulbaev]

Hello! I was exploring the new Linen API and found an error.
This is a toy example that differs from the tutorial in only two lines in the setup function (colab):

import functools
from typing import Any, Callable, Sequence, Optional
import numpy as np
import jax
from jax import lax, random, numpy as jnp
import flax
from flax.core import freeze, unfreeze
from flax import linen as nn

class ExplicitMLP(nn.Module):
  features: Sequence[int]

  def setup(self):
    # we automatically know what to do with lists, dicts of submodules
    self.layers = [nn.Dense(self.features[0])]
    self.layers += [nn.Dense(feat) for feat in self.features[1:]]
    # instead of self.layers = [nn.Dense(feat) for feat in self.features]
    # for single submodules, we would just write:
    # self.layer1 = nn.Dense(self, feat1)

  def __call__(self, inputs):
    x = inputs
    for i, lyr in enumerate(self.layers):
      x = lyr(x)
      if i != len(self.layers) - 1:
        x = nn.relu(x)
    return x

key1, key2 = random.split(random.PRNGKey(0), 2)
x = random.uniform(key1, (4,4))

model = ExplicitMLP(features=[3,4,5])
init_variables = model.init(key2, x)
y = model.apply(init_variables, x)

The error message:

ValueError: In setup, assign names of Modules via self.<name> and not using keyword argument name="<name>"

[avital]

Hi @Daulbaev -- there are two issues here:

1. The error message is misleading
2. Indeed, in setup, we attach submodules directly during assignment (via __setattr__). So Linen doesn't "see" the modification done with +=. There have been a few proposals for making this clearer, including: (a) making assigned values read only after assignment, (b) requiring the use of ModuleList, where we could then override __iadd__, (c) automatically converting assigned values to ModuleLists with that functionality.

[Daulbaev]

Got it. Thanks for the answer!

[avital]

So, for now, in your case you could do:

layers = [nn.Dense(self.features[0])]
layers += [nn.Dense(feat) for feat in self.features[1:]]
self.layers = layers