azula.guidance.tds

Twisted Posterior Sampling (TDS) internals.

References

Practical and Asymptotically Exact Conditional Sampling in Diffusion Models (Wu et al., 2023)

https://arxiv.org/abs/2306.17775

Classes

TDSSampler

Creates a TDS sampler.

Descriptions

class azula.guidance.tds.TDSSampler(denoiser, twist, **kwargs)

Creates a TDS sampler.

Parameters:

• denoiser (Denoiser) – A denoiser \(q_\phi(X \mid X_t)\).

• twist (Callable[[Tensor], Tensor]) – A twisting function $log p(y | hat{x}, t)$.

• kwargs – Keyword arguments passed to azula.sample.Sampler.

__call__(x, **kwargs)

Simulates the reverse process from \(t_T\) to \(t_0\).

Parameters:

• x (Tensor) – A set of \(K\) noisy tensors \(x_{t_T}\), with shape \((K, *)\).

• kwargs – Optional keyword arguments.

Returns:

The clean(er) tensors \(x_{t_0}\), with shape \((K, *)\).

Return type:

Tensor