VAE: Auto-Encoding Variational Bayes

class Encoder(nn.Module):
    """Map [B,Dx] inputs to [B,Dz] means and log variances (Appendix C.2)."""

    def __init__(self, input_dim, hidden_dim, latent_dim):
        super().__init__()
        self.hidden = nn.Linear(input_dim, hidden_dim)
        self.mu_head = nn.Linear(hidden_dim, latent_dim)
        self.logvar_head = nn.Linear(hidden_dim, latent_dim)

    def forward(self, x):
        """Return diagonal Gaussian parameters; logvar is log(sigma squared)."""
        check_matrix(x, "input", self.hidden.in_features)
        if x.dtype != self.hidden.weight.dtype:
            raise ValueError("input dtype must match encoder parameters")
        h = torch.tanh(self.hidden(x))
        return self.mu_head(h), self.logvar_head(h)

def reparameterize(mu, logvar, epsilon):
    """Use supplied standard-normal noise, preserving encoder gradients.

    Args:
        mu: Posterior means, shaped [B,Dz].
        logvar: Log variances, with the same shape and dtype as mu.
        epsilon: Independent standard-normal noise, shaped like mu.

    Returns:
        A differentiable [B,Dz] posterior sample.
    """
    check_pair(mu, logvar, "mu", "logvar")
    check_pair(mu, epsilon, "mu", "epsilon")
    std = torch.exp(0.5 * logvar)
    z = mu + std * epsilon
    check_matrix(z, "latent sample")
    return z

class Decoder(nn.Module):
    """Return [B,Dx] Bernoulli logits from [B,Dz] codes (Appendix C.1)."""

    def __init__(self, latent_dim, hidden_dim, input_dim):
        super().__init__()
        self.hidden = nn.Linear(latent_dim, hidden_dim)
        self.output = nn.Linear(hidden_dim, input_dim)

    def forward(self, z):
        """Return logits, rather than thresholded pixels or sampled images."""
        check_matrix(z, "latent sample", self.hidden.in_features)
        if z.dtype != self.hidden.weight.dtype:
            raise ValueError("latent dtype must match decoder parameters")
        return self.output(torch.tanh(self.hidden(z)))

def reconstruction_nll(logits, x):
    """Return feature-summed Bernoulli negative log likelihood per sample."""
    check_pair(logits, x, "logits", "input")
    if not bool(((x == 0) | (x == 1)).all()):
        raise ValueError("Bernoulli observations must be binary")
    element_nll = F.binary_cross_entropy_with_logits(
        logits, x, reduction="none"
    )
    return element_nll.sum(dim=1)

def analytic_kl(mu, logvar):
    """Sum diagonal Gaussian-to-standard-normal KL over latent dimensions."""
    check_pair(mu, logvar, "mu", "logvar")
    element_kl = 0.5 * (mu.square() + logvar.exp() - 1 - logvar)
    per_sample = element_kl.sum(dim=1)
    if not bool(torch.isfinite(per_sample).all()):
        raise ValueError("KL overflowed")
    return per_sample

def negative_elbo(logits, x, mu, logvar):
    """Minimize a one-sample negative ELBO; beta=1 and batch reduction=mean.

    Args:
        logits: Decoder Bernoulli logits, shaped [B,Dx].
        x: Binary observations, shaped like logits.
        mu: Posterior means, shaped [B,Dz].
        logvar: Posterior log variances, shaped like mu.

    Returns:
        Scalar loss, per-sample feature-summed NLL, and latent-summed KL.
    """
    nll = reconstruction_nll(logits, x)
    kl = analytic_kl(mu, logvar)
    if nll.shape != kl.shape or logits.dtype != mu.dtype:
        raise ValueError("likelihood and posterior must share batch and dtype")
    loss = (nll + kl).mean()
    if not bool(torch.isfinite(loss)):
        raise ValueError("objective overflowed")
    return loss, nll, kl

def training_step(encoder, decoder, x, epsilon, learning_rate=0.01):
    """Backpropagate both networks and take one SGD step with explicit noise.

    Args:
        encoder: A float CPU Encoder module.
        decoder: A compatible Decoder module.
        x: Binary observations, shaped [B,Dx].
        epsilon: Independent standard-normal noise, shaped [B,Dz].
        learning_rate: Positive finite SGD learning rate.

    Returns:
        Detached pre-update samples and objective terms, recorded gradients,
        and actual parameter-change evidence after the SGD update.
    """
    if not math.isfinite(learning_rate) or learning_rate <= 0:
        raise ValueError("learning_rate must be finite and positive")
    mu, logvar = encoder(x)
    z = reparameterize(mu, logvar, epsilon)
    logits = decoder(z)
    loss, nll, kl = negative_elbo(logits, x, mu, logvar)
    named = [("encoder." + k, p) for k, p in encoder.named_parameters()]
    named += [("decoder." + k, p) for k, p in decoder.named_parameters()]
    before = {name: p.detach().clone() for name, p in named}
    optimizer = torch.optim.SGD([p for _, p in named], lr=learning_rate)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    gradients = {name: p.grad.detach().clone() for name, p in named}
    if not all(bool(torch.isfinite(g).all()) for g in gradients.values()):
        raise ValueError("nonfinite gradients; parameters were not updated")
    optimizer.step()
    delta = sum(
        float((p.detach() - before[name]).square().sum()) for name, p in named
    )
    return {
        "mu": mu.detach(),
        "logvar": logvar.detach(),
        "z": z.detach(),
        "logits": logits.detach(),
        "objective": loss.detach(),
        "reconstruction_nll": nll.detach(),
        "kl": kl.detach(),
        "gradients": gradients,
        "parameter_delta_l2": math.sqrt(delta),
        "finite_parameters": all(
            bool(torch.isfinite(p).all()) for _, p in named
        ),
        "both_networks_updated": all(
            any(
                not torch.equal(p.detach(), before[name])
                for name, p in named
                if name.startswith(prefix)
            )
            for prefix in ("encoder.", "decoder.")
        ),
    }