Parameters in submodule/field not updated

[nalzok]

This is a follow-up from #2572.

After @cgarciae helped me figure out how to incorporate Distrax bijectors into a Flax module, I have a working VAE+Flow implementation here. However, the sampling result looks nearly the same regardless of whether I use a flow or not. After a closer inspection, I noticed that the parameters in the conditioners do not appear in the training state, as you can see by commenting out this line.

Basically, I have a Flax sub-sub-module (VAE (vae) -> Flow (vae.flow) -> Sequence[Dense] (vae.flow.conditioners)) whose parameters are not tracked/updated. How can I fix that?

[cgarciae]

Hey @nalzok, happy to help but can you post some reproducible snippet? 😅
I printed the variables in the previous issue and the conditioner parameters were there, intrigued as to how they are disappearing.

[nalzok]

Sure. You can clone the repo https://github.com/nalzok/vae-flow-flax and run make. There are just two files (which are very similar to the ones you saw yesterday) so probably not too overwhelming, but I can make a more minimal example once I get on my computer!

Ideally there should be things like conditioner_0 and conditioner_1 in the parameters, but I cannot find them.

[nalzok]

Hi @cgarciae, it turns out that you are right and the conditioner parameters are actually tracked/updated. I was mistaken since Flow is not initialized when flow_location is None, and I have for flow_location in (None, "prior", "posterior"): ..., and I only looked at the first iteration! 😅