Directly calling `module.setup()` should raise an exception

[andsteing]

flax.linen.Module.setup() should not be called directly because it needs a flax.linen.Scope to be set up properly.

Since #653 there is no more exception risen when a user inadvertently calls flax.linen.Module.setup() (though there probably will be error messages if the user tries to access the scope...)

[avital]

An incomplete thought below:

Curiously enough, this recent discussion made me think about this and wonder... For most Flax modules (that ultimately define parameters in a compact method), if you define submodules in setup but not parameters, I think it may be safe to call setup before a module is bound.

In fact, right now we have weird behavior where in Roland's example you can't run ResNet().backbone, even though it should probably work just fine (and still be an unbound module).

I'm not yet sure what the best answer is, but perhaps a workaround for Roland for now is actually to explicitly do:

resnet = ResNet()
resnet.setup()
resnet.backbone  # does this work?

[cgarciae]

Update

Currently @avital's example does not work:

resnet = ResNet()
resnet.setup()
resnet.backbone  # does this work?

You now get the following error:

AssertionError: Trying to register submodules on unbound scope.

This behaviour is documented in The Module lifecycle guide. We could go ahead and raise an error if setup is called on unbounded Module and add a message with some pointers into how to correctly extract submodules defined in setup.

Proposal

We could recommend the use of bind for this use-case, even a bind with empty variables could work:

resnet = ResNet()
backbone = resnet.bind({}).backbone.clone()

We could document this pattern in an existing guide or add a short guide about this topic of "Accessing Submodules".

Future ideas

If this pattern is sound we could make it a bit more user-friendly in the future via a .submodules that would automate the previous:

resnet = ResNet()
backbone = resnet.submodules.backbone