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).

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.reparameterizeDiag {α : Type} [TorchLean.Storage α] [Context α] {latent : Spec.Shape} (mu logvar eps : TorchLean.Tensor α latent) : TorchLean.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} [TorchLean.Storage α] [Context α] {latent : Spec.Shape} (mu logvar : TorchLean.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. For real-like scalars, the small-log-variance branch avoids subtracting one from a nearly unit exponential. A final nonnegative clamp removes negative roundoff; it is not an error bound. Complex inputs retain the original algebraic expression when the real-like guard fails.

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

A finite codebook for vector-quantized latent models.

• embedding : Fin numCodes → TorchLean.Tensor α latent

  Embedding vector for each code index.

def Generative.Latent.quantizeAt {α : Type} [TorchLean.Storage α] {numCodes : ℕ} {latent : Spec.Shape} (book : Codebook α numCodes latent) (idx : Fin numCodes) : TorchLean.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.