Wrong parameter names when nesting Modules within flax transformations about flax HOT 3 OPEN

Hey @PhilipVinc, what is happening is that nn.Dense(features=1) is being attached to VNet during VNet.__post_init__, as its found inside the some_args field. We eagerly bind submdoules found in fields to avoid race conditions and make parenting more deterministic.

from flax.

Ah, I see. Thanks for the answer

Is there some alternative workaround ? What I want to achieve Is actually to be able to build VNet by just passing the underlying module to be lifted/vmapped.

However the lifted flax.nn.vmap wants the module (constructor) and so I'm obliged to pass the arguments, which in this case happen to also contain a module.

ideally, the usage I'd like is for this to work:

sub_subnet =  nn.Dense(features=1)
sub_net = Net(subnet= sub_subnet)
net = VNet(subneta= sub_net)

A bit like I'm able to pass an already constructed module to a Module, I would also like to pass an already constructed module (with submodules) to VNet...

from flax.