Stochastic Layers#

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 (optionally) samples S latent draws via the reparameterization trick during training.

The build step infers F and lazily constructs the linear layers mu and log_var.

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

Compute parameters and (optionally) samples from the Gaussian posterior.

During training, this method returns S Monte Carlo samples using the reparameterization trick:

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

During evaluation (model.eval()), deterministic output is returned with z equal to the repeated mean.

Parameters:

x (torch.Tensor) – Input tensor of shape [B, F].
S (int, optional) – Number of Monte Carlo samples to generate. Must be >= 1. Defaults to 1 (single sample).

Returns:

(z, mu, log_var), where:

Return type:

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

Raises:

ValueError – If S < 1.

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

Compute posterior parameters without drawing samples.