Latent-variable generative helpers

This module contains the shared spec-layer vocabulary for latent generative models:

• continuous latent variables, as used in variational autoencoders (VAEs);
• discrete codebook latents, as used in vector-quantized VAEs (VQ-VAEs); and
• small total scalar/tensor helpers that keep model files focused on architecture.

The definitions are intentionally model-agnostic. A VAE, VQ-VAE, latent diffusion model, or normalizing-flow model can all reuse these primitives without committing to a particular backbone.

References:

• Kingma and Welling (2014), "Auto-Encoding Variational Bayes" (VAE).
• Rezende, Mohamed, and Wierstra (2014), "Stochastic Backpropagation and Approximate Inference".
• van den Oord, Vinyals, and Kavukcuoglu (2017), "Neural Discrete Representation Learning" (VQ-VAE).

def Generative.Latent.expTensor {α : Type} [Context α] {s : Spec.Shape} (x : Spec.Tensor α s) : Spec.Tensor α s

Elementwise exponential, useful for log-variance parameterizations.

def Generative.Latent.halfTensor {α : Type} [Context α] {s : Spec.Shape} (x : Spec.Tensor α s) : Spec.Tensor α s

Elementwise $\tfrac12x$, written as a tensor helper to make VAE equations readable.

def Generative.Latent.reparameterizeDiag {α : Type} [Context α] {latent : Spec.Shape} (mu logvar eps : Spec.Tensor α latent) : Spec.Tensor α latent

Diagonal-Gaussian reparameterization:

$$ z=\mu+\exp\!\left(\tfrac12\log\sigma^2\right)\odot\varepsilon. $$

This is the spec-level form of the VAE reparameterization trick. The noise $\varepsilon$ is explicit, so the function stays pure and deterministic; runtime examples can supply deterministic or random noise.

def Generative.Latent.diagonalGaussianKlToStandard {α : Type} [Context α] {latent : Spec.Shape} (mu logvar : Spec.Tensor α latent) : α

Mean KL term for a diagonal Gaussian posterior against a standard normal prior.

For each latent coordinate:

$$ D_{\mathrm{KL}}\!\left(\mathcal{N}(\mu,\sigma^2)\,\middle\|\,\mathcal{N}(0,1)\right) =\tfrac12\left(\exp(\log\sigma^2)+\mu^2-1-\log\sigma^2\right). $$

We return the mean across the latent shape, matching TorchLean's existing loss convention.

structure Generative.Latent.Codebook (α : Type) (numCodes : ℕ) (latent : Spec.Shape) [Context α] : Type

A finite codebook for vector-quantized latent models.

• embedding : Fin numCodes → Spec.Tensor α latent

  Embedding vector for each code index.

def Generative.Latent.quantizeAt {α : Type} [Context α] {numCodes : ℕ} {latent : Spec.Shape} (book : Codebook α numCodes latent) (idx : Fin numCodes) : Spec.Tensor α latent

Quantize by an explicit code index.

Nearest-neighbor lookup is usually how VQ-VAE chooses this index during execution. The spec keeps the index explicit so proofs and verifiers can reason about a fixed code assignment without depending on an argmin/tie-breaking policy.