Stochastic Layers#

Sampling layer of VAE.

A VAE creates its FullyFactorizedGaussian as vae.sampling_layer. To use the class directly, import it from this module; it is not exported by pyautoencoder.variational:

from pyautoencoder.variational.stochastic_layers import FullyFactorizedGaussian
class pyautoencoder.variational.stochastic_layers.FullyFactorizedGaussian(latent_dim: int)#

Bases: Module

Gaussian posterior head producing a fully factorized \(q(z \mid x)\).

Given input features x of shape [B, F], this module produces the parameters of a diagonal Gaussian posterior,

\[q(z \mid x) = \mathcal{N}(z \mid \mu(x), \operatorname{diag}(\sigma(x)^2)),\]

and either S latent samples drawn via the reparameterization trick (sample=True, the default) or the posterior mean (sample=False). The choice is explicit and does not depend on train/eval mode.

The linear layers mu and log_var are torch.nn.LazyLinear: their input size F is inferred on the first forward pass, for example the one run by VAE.build.

forward(x: Tensor, S: int = 1, sample: bool = True) → tuple[Tensor, Tensor, Tensor]#

Compute the posterior parameters and the latent codes z.

With sample=True (default), z holds S Monte Carlo samples drawn with the reparameterization trick (see reparametrize()):

\[z^{(s)} = \mu + \sigma \odot \epsilon^{(s)}, \qquad \epsilon^{(s)} \sim \mathcal{N}(0, I).\]

With sample=False, z is the posterior mean mu, with a single sample dimension (S must be 1).

The behavior does not depend on train/eval mode.

Parameters:

  • x (torch.Tensor) – Input tensor of shape [B, F].

  • S (int, optional) – Number of Monte Carlo samples to generate. Must be >= 1, and 1 when sample=False. Defaults to 1.

  • sample (bool, optional) – Whether to sample z from \(q(z \mid x)\) (True) or to return its mean (False). Defaults to True.

Returns:

(z, mu, log_var), where:

  • z – latent codes, shape [B, S, latent_dim] ([B, 1, latent_dim] when sample=False).

  • mu – mean of \(q(z \mid x)\), shape [B, latent_dim].

  • log_var – log-variance of \(q(z \mid x)\), shape [B, latent_dim].

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Raises:

ValueError – If S < 1, if S != 1 with sample=False, or if x is not of shape [B, F].

get_params(x: Tensor) → tuple[Tensor, Tensor]#

Compute posterior parameters without drawing samples.

Parameters:

x (torch.Tensor) – Input tensor of shape [B, F].

Returns:

(mu, log_var), each of shape [B, latent_dim].

Return type:

tuple[torch.Tensor, torch.Tensor]

Raises:

ValueError – If x is not of shape [B, F].

reparametrize(mu: Tensor, log_var: Tensor, S: int = 1) → Tensor#

Draw S latent samples via the reparameterization trick.

\[z^{(s)} = \mu + \sigma \odot \epsilon^{(s)}, \qquad \epsilon^{(s)} \sim \mathcal{N}(0, I).\]

This method always samples. It is differentiable with respect to mu and log_var, so it can be used on posterior parameters computed elsewhere (for example the adapted posteriors of AdaGVAE).

Parameters:

  • mu (torch.Tensor) – Mean of the posterior, shape [B, D_z].

  • log_var (torch.Tensor) – Log-variance of the posterior, shape [B, D_z].

  • S (int, optional) – Number of samples to draw. Must be >= 1. Defaults to 1.

Returns:

Sampled latent codes of shape [B, S, D_z].

Return type:

torch.Tensor

Raises:

ValueError – If S < 1, or if mu and log_var are not both of shape [B, D_z].