Declare linen submodule in dataclass-style?

Hi all,

First of all, thank you for your work on flax. This is not an issue but a question.

As a pytorch user, sometimes I want to create a submodule, then pass it as an attribute to another module. This is quite common because people want to have flexibility about how to pick submodules and assemble them into the model. For instance:

dummy = Dummy()
bigdummy = BigDummy(dummy = dummy)
init_vars = bigdummy.init(rngkey, x)

the modules are defined as

class Dummy(nn.Module):
  def setup(self):
    self.dense = nn.Dense(4)

  def __call__(self, x):
    return self.dense(x)

class BigDummy(nn.Module):
  dummy: Dummy

  def setup(self):
    self.proj = nn.Dense(6)

  def __call__(self, x):
    return self.proj(self.dummy(x))

However, this fails at bigdummy.init for attribute access bigdummy.dummy because it doesn’t run dummy.setup. I guess this has to do with how flax overrides __setattr__. So I’m wondering is there any way that flax supports this use case?

Thank you in advance!

Environment: flax 0.3

[levskaya]

@chris-tng @avital this has been a commonly cited pain-point, and I’ve tried to remedy it with PR #864 - I believe your example should work fine with this change. (For the record I don’t think this is actually so related to #686 - that’s more of an interactive use-case question of inspecting instantiated parameters, etc. whereas this is simply wanting to define modules made from simpler ones at the top-level following the common pytorch pattern.)

[avital]

Let’s not close this issue just yet – I’d like for whoever works on #686 (probably me or @jheek) to keep this use-case in mind.

I agree this is suboptimal. It may be solvable separate from #686.