Variational Autoencoder

"""A VAE forward pass and its ELBO - the two terms, written out.

Reconstruction pulls the latent toward encoding the input; the KL pulls it
toward the prior. The model that is useful sits between them.
"""

import math
import random


def encode(x, w_mu, w_logvar):
    """Returns the posterior mean and log-variance for one input."""
    mu = [sum(w * v for w, v in zip(row, x)) for row in w_mu]
    logvar = [sum(w * v for w, v in zip(row, x)) for row in w_logvar]
    return mu, logvar


def reparameterize(mu, logvar):
    """z = mu + sigma * eps, with eps sampled OUTSIDE the computation.

    This is the whole trick: the randomness becomes an input rather than an
    operation, so a gradient can flow through to mu and logvar.
    """
    z = []
    for m, lv in zip(mu, logvar):
        sigma = math.exp(0.5 * lv)
        eps = random.gauss(0.0, 1.0)
        z.append(m + sigma * eps)
    return z


def decode(z, w_out):
    return [sum(w * v for w, v in zip(row, z)) for row in w_out]


def elbo(x, x_hat, mu, logvar):
    """The bound, as two explicit terms. Returned to be MAXIMIZED."""
    # Gaussian log-likelihood, up to a constant: negative squared error.
    reconstruction = -sum((a - b) ** 2 for a, b in zip(x, x_hat))

    # Closed form for KL(N(mu, sigma) || N(0, I)) - no sampling needed.
    kl = 0.5 * sum(
        math.exp(lv) + m * m - 1.0 - lv for m, lv in zip(mu, logvar)
    )

    return reconstruction - kl

"""VAE - typed, validated, with the KL warm-up that prevents collapse."""

from dataclasses import dataclass

import torch
from torch import Tensor, nn


@dataclass(frozen=True)
class VaeConfig:
    input_dim: int
    latent_dim: int = 32
    hidden_dim: int = 256
    beta: float = 1.0


class Vae(nn.Module):
    def __init__(self, config: VaeConfig) -> None:
        super().__init__()
        if config.latent_dim < 1:
            raise ValueError(f"latent_dim must be >= 1, got {config.latent_dim}")
        if config.beta < 0:
            raise ValueError(f"beta must be non-negative, got {config.beta}")

        self.config = config
        self.encoder = nn.Sequential(
            nn.Linear(config.input_dim, config.hidden_dim),
            nn.ReLU(),
        )
        # Two heads, not one: the posterior needs a mean and a spread.
        self.to_mu = nn.Linear(config.hidden_dim, config.latent_dim)
        self.to_logvar = nn.Linear(config.hidden_dim, config.latent_dim)
        self.decoder = nn.Sequential(
            nn.Linear(config.latent_dim, config.hidden_dim),
            nn.ReLU(),
            nn.Linear(config.hidden_dim, config.input_dim),
        )

    def encode(self, x: Tensor) -> tuple[Tensor, Tensor]:
        h = self.encoder(x)
        # log-variance rather than variance, so the output is unconstrained
        # and exp() guarantees positivity without a clamp.
        return self.to_mu(h), self.to_logvar(h)

    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        if x.dim() != 2:
            raise ValueError(f"expected (batch, features), got {tuple(x.shape)}")
        mu, logvar = self.encode(x)
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)
        return self.decoder(z), mu, logvar


def elbo_loss(
    x: Tensor, x_hat: Tensor, mu: Tensor, logvar: Tensor, beta: float
) -> tuple[Tensor, Tensor]:
    """Returns (negative ELBO to minimize, KL term for monitoring).

    The KL is returned separately on purpose: posterior collapse is invisible
    in the total, and the total is what people watch.
    """
    reconstruction = nn.functional.mse_loss(x_hat, x, reduction="sum") / x.size(0)
    kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / x.size(0)
    return reconstruction + beta * kl, kl


def kl_weight(epoch: int, warmup_epochs: int, beta: float) -> float:
    """Anneal the KL from 0 to beta. Without this the KL wins early, the
    encoder collapses to the prior, and it never recovers."""
    if warmup_epochs <= 0:
        return beta
    return beta * min(1.0, epoch / warmup_epochs)

"""VAE - fused KL, multi-sample estimator, one pass over the batch."""

import torch
from torch import Tensor, nn


def elbo_fused(x: Tensor, x_hat: Tensor, mu: Tensor, logvar: Tensor, beta: float) -> Tensor:
    """Both ELBO terms in a single fused reduction.

    The naive form builds five intermediate tensors (logvar.exp(), mu.pow(2),
    the sum, the negation, the scale). torch.addcmul folds the square and the
    accumulate into one traversal, and the whole expression reduces in one pass.
    """
    batch = x.size(0)
    recon = (x_hat - x).pow_(2).sum() / batch

    # -0.5 * sum(1 + logvar - mu^2 - exp(logvar)), computed without temporaries.
    kl_terms = torch.addcmul(logvar.exp() - logvar - 1.0, mu, mu)
    kl = 0.5 * kl_terms.sum() / batch

    return recon + beta * kl


def importance_weighted_elbo(model: nn.Module, x: Tensor, num_samples: int = 8) -> Tensor:
    """A tighter bound from K latent samples, for anomaly scoring.

    Expanding the batch to (K*B, D) evaluates all K samples in ONE decoder
    pass rather than K sequential ones - the difference between an 8x cost and
    roughly a 1.2x one, since the decoder is compute-bound and under-occupied
    at typical batch sizes.
    """
    batch, dim = x.shape
    mu, logvar = model.encode(x)

    # Contiguous expansion so the decoder sees one dense batch, not a strided view.
    mu_rep = mu.repeat_interleave(num_samples, dim=0)
    logvar_rep = logvar.repeat_interleave(num_samples, dim=0)
    x_rep = x.repeat_interleave(num_samples, dim=0)

    z = mu_rep + torch.exp(0.5 * logvar_rep) * torch.randn_like(mu_rep)
    x_hat = model.decoder(z)

    log_w = -(x_hat - x_rep).pow(2).sum(dim=1)
    log_w = log_w.view(batch, num_samples)

    # logsumexp in float32 regardless of autocast dtype: the whole point of
    # the estimator is the tail, and the tail is where float16 underflows.
    return torch.logsumexp(log_w.float(), dim=1) - torch.log(
        torch.tensor(float(num_samples))
    )