azula.guidance.mmps¶

Moment Matching Posterior Sampling (MMPS) internals.

References

Learning Diffusion Priors from Observations by Expectation Maximization (Rozet et al., 2024)

Classes¶

Creates a MMPS denoiser module.

class azula.guidance.mmps.MMPSDenoiser(denoiser, y, A, cov_y, solver='gmres', iterations=1)¶

Creates a MMPS denoiser module.

Parameters:

denoiser (Denoiser) – A denoiser \(q_\phi(X \mid X_t)\).
y (Tensor) – An observation \(y \sim \mathcal{N}(A(x), \Sigma_y)\), with shape \((*, D)\).
A (Callable[[Tensor], Tensor]) – The forward operator \(x \mapsto A(x)\).
cov_y (Tensor | Covariance) – The noise covariance \(\Sigma_y\). If cov_y is a tensor, it is assumed to be the variance \(\sigma_y^2\) and \(\Sigma_y = \mathrm{diag}(\sigma_y^2)\).
solver (str) – The linear solver name ("cg" or "gmres").
iterations (int) – The number of solver iterations.