Reverse DDPM sampler (spec layer)

This file defines a standard ε-prediction reverse sampler step for discrete VP/DDPM schedules.

We expose:

• ddpmStep: one reverse step $x_t\to x_{t-1}$ with explicit noise input $z$,
• ddpmSample: run all $T$ reverse steps, given a noise stream $z_0,\ldots,z_{T-1}$.

We keep everything scalar-polymorphic (Context α). The intended use is:

• execute with Float, Float32, or the configured binary32 type ExecFloat.Binary 8 23;
• run in CPU software at a chosen precision with FloatLib.Floats.ExecFloat.Binary; and
• reuse the definitions with ℝ or noncomputable FloatLib.Floats.Formats.Flocq.NF in proofs.

References (informal pointers):

• Ho, Jain, Abbeel (2020), DDPM, Algorithm 2 (reverse process).

def Generative.Diffusion.x0PredFromEps {α : Type} [TorchLean.Storage α] [Context α] {T : ℕ} {s : Spec.Shape} (sched : VPSchedule α T) (x_t epsHat : TorchLean.Tensor α s) (t : Fin (T + 1)) : TorchLean.Tensor α s

Reconstruct $x_0$ from $x_t$ and an already computed noise prediction.

$$ x_0=\frac{x_t-\sqrt{1-\bar\alpha_t}\,\hat\varepsilon} {\sqrt{\bar\alpha_t}}. $$

safeDiv adds the context's epsilon to the denominator, as in x0Pred. Passing the prediction explicitly lets DDIM use the same tensor for this reconstruction and the direction toward the previous sample.

def Generative.Diffusion.x0Pred {α : Type} [TorchLean.Storage α] [Context α] {T : ℕ} {s : Spec.Shape} (sched : VPSchedule α T) (model : EpsModel α s) (x_t : TorchLean.Tensor α s) (t : Fin (T + 1)) : TorchLean.Tensor α s

Predict $x_0$ by evaluating the denoiser at time t / T and applying x0PredFromEps.

Call x0PredFromEps directly when another part of the sampler already needs the same denoiser output. Both entry points use the same coefficient arithmetic and additive epsilon shift.

def Generative.Diffusion.ddpmStep {α : Type} [TorchLean.Storage α] [Context α] {T : ℕ} {s : Spec.Shape} (sched : VPSchedule α T) (model : EpsModel α s) (k : Fin T) (x_t z : TorchLean.Tensor α s) : TorchLean.Tensor α s

One reverse DDPM step $x_t\to x_{t-1}$ with explicit noise $z$ (intended as $\mathcal{N}(0,I)$).

We index reverse steps by k : Fin T corresponding to the transition $t=k+1\to k$.

Implementation details:

• time embedding passed to the model is $t/T$ (see VPSchedule.timeOfIndex).
• coefficient divisions add Context.defaultEpsilon to the denominator. The shift is part of this sampler's formula and does not guarantee nonzero denominators or finite outputs.

def Generative.Diffusion.ddpmSample {α : Type} [TorchLean.Storage α] [Context α] {T : ℕ} {s : Spec.Shape} (sched : VPSchedule α T) (model : EpsModel α s) (x_T : TorchLean.Tensor α s) (noise : Fin T → TorchLean.Tensor α s) : TorchLean.Tensor α s

Run the full reverse DDPM sampler for T steps.

Inputs:

• $x_T$: starting state (typically standard normal noise),
• noise: per-step noise stream z_k for k = 0..T-1.

Output:

• the terminal sample $x_0$.

Order note:

• noise (T-1) is used first (for the step $T\to T-1$),
• noise 0 is used last (for the step $1\to0$).