How does Module.setup() in Flax relate to build() in TF2 Keras?

[marcvanzee]

(Original question from Han Zhang)

The TF2 Keras API has a build function (documentation). This function gets called only once, and returns the input shape which can be used to create some layers depend on the input shape.

linen.Module has setup(), according to the doc:

setup() is called once lazily on a module instance when a module
is bound, immediately before any other methods like ``__call__`` are
invoked, or before a ``setup``-defined attribute on `self` is accessed.

I originally thought it would be similar to build function in TF2, however, I could not get the input shape using setup(). Why is this the case?

[marcvanzee]

Answer from @levskaya:

the shape parameters to setup are meant to come from the dataclass attributes. Since by default submodules are usually lazy themselves, they can init their parameters once __call__'d with access to more shape info. If you need to create a -variable- in setup then you have to provide the shape info via attributes.