AbstractVAE methods

[vitskvara]

I have an issue with methods such as

encoder_mean(m::AbstractVAE, x::AbstractArray) = m.encoder(x)

This is not gonna work with a classical VAE where the last layer of encoder contains both means and sigmas. Maybe change this to

encoder_mean(m::AbstractVAE, x::AbstractArray) = mean(m,m.encoder(x))

or something similar? That would however require that the model contains some label with what kind of VAEit is. This is also ttrue for decoder, where you can have fixed/scalar/diagonal sigmas. Or is there some advantage of using the TrackedArray sigmas separated as in your case?

[nmheim]

Yes, good point. I have been wondering how to implement VAE with fixed/loose sigmas elegantly. Maybe it would be easiest to have distinct subtypes for each variant instead of labels? not sure how much extra code that would cause....

regarding the seperated sigma, I just want to use a shared sigma for all data points.

[vitskvara]

Yes, but this could lead to a lot of different subtypes. However if we keep some methods mutual than it could still work.

I see, however this should also cover the case when you don't want to share sigma and instead have sigma(x).

I think we will have to trade between generality of the code and its applicability to the different applications that we have.

[vitskvara]

The answer to the first question is the same as to whether we should have a separate type for VAE with KL divergence and VAE with Wasserstein, where the prior can be basically anything, or try to fit everything into one type with different loss functions.

[smidl]

I do not think so. We can define elbo() and wasserstein() which one is used in training is up to the user.

The number of cases for variance of encoder is not too big. We can have fixed, scalar(x), diagonal(x) and that is pretty much it. The same for sigma_e. So, I would vote for a general model with this kind of switch. I admit that it may lead to little mess in the attributes -- we need scalar attibute for the constant which will be unused when variance is from the network. But it is not a big deal, right?