dew.sampling

The reverse process for diffusion, and decoding for language models.

Name	Summary
CFG	Interval-limited classifier-free guidance (Kynkaanniemi et al.
DDIM	DDIM (Song et al.
DDPM	Exact ancestral sampler for the reverse diffusion SDE.
DEIS	DEIS (Zhang and Chen 2023, arXiv 2204.13902) in its log-rho multistep form, Diffusers 0.34.0’s DEISMultistepScheduler: the exponential integrator of eps with the polynomial-in-log(rho) interpolation of the last outputs, rho = sigma / alpha, integrated in closed form over the step.
KDPM2	k-diffusion’s DPM-Solver-2 (sample_dpm_2), the update of Diffusers 0.34.0’s KDPM2DiscreteScheduler, and with ancestral its sample_dpm_2_ancestral and KDPM2AncestralDiscreteScheduler.
LMS	Linear multistep over dx/dsigma = (x - x_0) / sigma, k-diffusion’s sample_lms and Diffusers 0.34.0’s LMSDiscreteScheduler: the last order derivatives interpolated by the Lagrange polynomial through their sigmas and integrated over the step, in closed form where Diffusers quadratures; the order grows with the history.
PNDM	PNDM (Liu et al.
RK4	Classical Runge-Kutta over dx/dsigma = eps, on a variance exploding schedule; the stages at half steps read the model at the time the schedule maps that sigma back to.
TCD	Trajectory consistency sampling (Zheng et al.
Beam	Deterministic beam search over one shared prefill.
Consistency	Multistep consistency sampling (Song et al.
DPMSolverMultistep	DPM-Solver (Lu et al.
DPMSolverSDE	Diffusers 0.34.0’s DPMSolverSDEScheduler, k-diffusion’s sample_dpmpp_sde midpoint solver over a Brownian tree.
DPMSolverSinglestep	Diffusers 0.34.0’s grouped DPM-Solver updates from each group’s anchor.
Euler	The DDIM update written as an Euler step of the probability flow ODE.
EulerAncestral	Euler with the ancestral noise injection of k-diffusion (get_ancestral_step, eta 1).
FlowSDE	Flow-GRPO’s Euler-Maruyama solver on a rectified-flow Process.
FlowTrajectory	A reverse trajectory, with batch-major states and joint log densities.
GaussianTransition	An isotropic transition with one variance per batch row.
Generation	Prompt plus padded continuation, and response-aligned likelihoods.
Grammar	A token automaton on the device.
Heun	Heun’s second order method (Karras et al.
LogitsTransform	A pure [rows, vocab] score rewrite, applied before the draw. Documented in dew.sampling.decoding.
MultiStepDPM	A third order multistep integrator of dx/dsigma = eps on a variance exploding schedule, from finite differences of the last three eps.
Sample	Draw every row independently, one token per step.
Sampling	Token selection and termination.
Solver	A step of a sampler, and whatever it carries between steps.
Speculative	Draft with the model’s prediction depths, verify with the model itself.
StepState	What a transform or a criterion sees at one decode step. Documented in dew.sampling.decoding.
Stopping	A pure per-row finish test over the tokens a step just drew. Documented in dew.sampling.decoding.
Strategy	The device loop of one generation request.
TextToImage	pipe(prompts, seed=0) or pipe(prompts, steps=40, guidance=4.0, sampler=samplers.Heun(), key=key).
UniPC	UniPC (Zhao et al.
flow_transition	Euler-Maruyama over the physical noise rate, with 0 <= sigma_next <= sigma <= 1.
generate	Generate from numeric model inputs, with an array shorthand for text.
sample	steps points from T to 0: a solver step across each interval, then the model’s clean prediction at the last point.
sample_trajectory	Record FlowSDE transitions over the same time grid and keys as sample.

CFG

dataclass source

class CFG(
    scale: float,
    interval: tuple[float, float] = (0.0, 1.0),
    rescale: float = 0.0,
)

Interval-limited classifier-free guidance (Kynkaanniemi et al. 2024).

The guided prediction is uncond + scale (cond - uncond). Guidance hurts at high noise and buys nothing at low noise, so outside interval the scale drops to 1, which is exactly the plain conditional prediction. The interval is in trajectory progress, 0 at pure noise and 1 at the clean sample; the default covers all of it.

rescale is the guidance rescaling of Lin et al. 2023 (“Common Diffusion Noise Schedules and Sample Steps are Flawed”, section 3.4), Diffusers’ guidance_rescale: the guided output is rescaled to the per-sample standard deviation of the conditional one and mixed back at that weight, so 0 leaves the guided output alone and 1 takes the rescaled one. The standard deviation is over everything but the batch axis, with the unbiased correction the reference’s Tensor.std applies.

Guidance combines the model’s raw outputs and lets the denoiser convert once, so a source’s clipping, dynamic thresholding or consistency boundary sees the guided output rather than each branch separately.

DDIM

dataclass source

class DDIM(eta: float = 0.0)

DDIM (Song et al. 2021); eta is the stochasticity, 0 deterministic and 1 DDPM-like.

Diffusers 0.34.0’s DDIMScheduler limits the clean prediction under clip_sample or thresholding and keeps the model’s own output as its epsilon, so the direction term is the unlimited one. That pairing is SourceLimitedPrediction’s, in the process’s conversion.

DDIM.init

def init(x, times, process, *, key)

DDIM.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DDPM

dataclass source

class DDPM(variance: Literal['small', 'large'] = 'small')

Exact ancestral sampler for the reverse diffusion SDE.

One step draws from the forward posterior q(x_s | x_t, x_0) for x_t = alpha_t x_0 + sigma_t eps, written in signal and noise rates so it holds for any schedule and any step stride. The posterior mean is alpha_s x_0 + alpha_t sigma_s^2 / (alpha_s sigma_t) eps and its variance is sigma_s^2 (1 - alpha_t^2 sigma_s^2 / (alpha_s^2 sigma_t^2)).

variance is which of Diffusers 0.34.0’s fixed DDPMScheduler posterior variances the draw takes. "small" is that posterior’s own, written in rates and so defined on any schedule. "large" is the forward step’s beta, 1 - alpha_t^2 / alpha_s^2, the wider Glide choice: that is a variance-preserving statement, and it is zero wherever alpha is one, so a variance-exploding grid is refused rather than sampled without noise.

Neither draws on the step whose own time is the schedule’s zero: x_t is the least noised state the schedule holds there, and the source gates its draw on that time the same way. Elsewhere the wide variance at a terminal alpha of one is exactly sigma_t, which is what the source’s own current_beta_t reduces to.

DDPM.init

def init(x, times, process, *, key)

DDPM.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DEIS

dataclass source

class DEIS(order: int = 2, lower_order_final: bool = True)

DEIS (Zhang and Chen 2023, arXiv 2204.13902) in its log-rho multistep form, Diffusers 0.34.0’s DEISMultistepScheduler: the exponential integrator of eps with the polynomial-in-log(rho) interpolation of the last outputs, rho = sigma / alpha, integrated in closed form over the step. The first order is DPM-Solver’s. Orders grow with the history; lower_order_final is the same short-walk taper as DPMSolverMultistep’s. At sigma=0 the integrated log-rho basis retains its history terms. An alpha=0 source contributes a node at infinite rho; that node’s weight vanishes in subsequent finite-interval integrals.

DEIS.init

def init(x, times, process, *, key)

DEIS.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

KDPM2

dataclass source

class KDPM2(ancestral: bool = False)

k-diffusion’s DPM-Solver-2 (sample_dpm_2), the update of Diffusers 0.34.0’s KDPM2DiscreteScheduler, and with ancestral its sample_dpm_2_ancestral and KDPM2AncestralDiscreteScheduler.

An Euler step to the geometric midpoint of sigma_t and the target level, the model read there, and the step from x taken with that midpoint derivative. The target is sigma_s, or under ancestral the sigma_down of k-diffusion’s ancestral step with sigma_up of fresh noise added after. The midpoint’s time comes from the schedule’s t_of_sigma, so this integrates a GeneralizedNoiseScheduler.

KDPM2.init

def init(x, times, process, *, key)

KDPM2.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

LMS

dataclass source

class LMS(order: int = 4)

Linear multistep over dx/dsigma = (x - x_0) / sigma, k-diffusion’s sample_lms and Diffusers 0.34.0’s LMSDiscreteScheduler: the last order derivatives interpolated by the Lagrange polynomial through their sigmas and integrated over the step, in closed form where Diffusers quadratures; the order grows with the history. Integrates a GeneralizedNoiseScheduler.

LMS.init

def init(x, times, process, *, key)

LMS.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

PNDM

dataclass source

class PNDM(skip_prk_steps: bool = False)

PNDM (Liu et al. 2022, arXiv 2202.09778), Diffusers 0.34.0’s PNDMScheduler: Adams-Bashforth over eps with DDIM as the transfer, fourth order once four outputs are in hand. The warmup is the paper’s, three pseudo Runge-Kutta steps of four model evaluations each (the stages at the interval’s midpoint and end), which seed the history with each step’s first eps; under skip_prk_steps, the PLMS form Stable Diffusion runs, the first step is a predictor-corrector pair (an Euler step, eps re-read at its end, the step retaken with the mean) and the orders grow from there. The schedule owns the transfer stride: ordinary native grids use their adjacent interval, while published integer grids retain their fixed training stride. Transfers cannot start at alpha = 0.

PNDM.init

def init(x, times, process, *, key)

PNDM.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

RK4

dataclass source

class RK4()

Classical Runge-Kutta over dx/dsigma = eps, on a variance exploding schedule; the stages at half steps read the model at the time the schedule maps that sigma back to.

RK4.init

def init(x, times, process, *, key)

RK4.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

TCD

dataclass source

class TCD(eta: float = 0.3)

Trajectory consistency sampling (Zheng et al. 2024, arXiv 2402.19159), Diffusers 0.34.0’s TCDScheduler: the deterministic DDIM step lands at (1 - eta) t_next, and the forward process noises it up to t_next with fresh noise, the paper’s gamma-sampling with gamma = eta. At eta 0 it is DDIM; at eta 1 the step goes through the clean end of the schedule. A tabulated schedule reads the intermediate time at its truncated index, as the reference floors it. The last step of a walk lands on the grid’s end itself and takes no noise.

TCD.init

def init(x, times, process, *, key)

TCD.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

Beam

dataclass source

class Beam(
    width: int = struct.field(pytree_node=False, default=1),
    length_penalty: float = struct.field(pytree_node=False, default=1.0),
    early_stopping: bool | str = struct.field(pytree_node=False, default=False),
    stop_ids: int = struct.field(pytree_node=False, default=1),
)

Deterministic beam search over one shared prefill.

The bookkeeping is _beam_search in Transformers 5.16.1. A step scores every live beam’s continuations, keeps the best (1 + stop_ids) * width of them so width live beams always remain, moves the ones a criterion ended into the completed set with their score divided by their generated length raised to length_penalty, and continues with the rest. early_stopping follows the reference’s three settings: False estimates the best score still reachable from the current length, True also stops recording once every beam is completed, and “never” estimates from the whole budget when the penalty rewards length.

The prompt is prefilled once and its cache row is copied into width rows; every step reparents those rows through DecodeOps.reindex, so a branched beam decodes exactly like a separately selected prefix. Parameters are never mapped.

n is how many completed beams to return, not the search width, and n > width is an error. A selected path is a search result rather than a draw, so its behaviour log probability is zero; the raw log probabilities stay the model’s own for the tokens on the path.

keep: int

Continuations a step keeps, as beams_to_keep upstream.

Consistency

dataclass source

class Consistency()

Multistep consistency sampling (Song et al. 2023, Algorithm 1), the update of Diffusers 0.34.0’s LCMScheduler: the clean prediction is noised again to the next level with fresh noise, x_s = alpha_s x_0 + sigma_s z, and the last step keeps x_0 as it is. A latent consistency model’s x_0 is the consistency function’s output, which ConsistencyBoundary reads out of the model’s prediction.

Consistency.init

def init(x, times, process, *, key)

Consistency.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DPMSolverMultistep

dataclass source

class DPMSolverMultistep(
    order: int = 2,
    algorithm: Algorithm = 'dpmsolver++',
    solver_type: Literal['midpoint', 'heun'] = 'midpoint',
    lower_order_final: bool = True,
    euler_at_final: bool = False,
)

DPM-Solver (Lu et al. 2022, arXiv 2206.00927) and DPM-Solver++ (arXiv 2211.01095) as multistep integrators in lambda = log(alpha) - log(sigma): the four algorithms, three orders and two second-order forms of Diffusers 0.34.0’s DPMSolverMultistepScheduler, with its defaults.

dpmsolver++ and sde-dpmsolver++ integrate the clean prediction, the other two eps; the sde- forms add the noise term of the SDE solver. With h = lambda_t - lambda_s0 over the step and D0, D1, D2 the finite differences of the last outputs in lambda, the deterministic dpmsolver++ step is

x_t = (sigma_t / sigma_s0) x - alpha_t (e^-h - 1) D0 + c D1 + c_2 D2

and _dpm_terms holds each algorithm’s coefficients as Diffusers writes them. The first step has no history and is first order, the second at most second. lower_order_final is Diffusers’ rule verbatim, which acts only in a walk under 15 steps: first order on the last step and at most second on the one before. euler_at_final makes the last step first order whatever the length. A zero target forces the clean limit for deterministic and ++ updates. The non-++ SDE first-order limit also retains alpha_tsigma_s/alpha_s(noise-eps); it requires a finite source alpha and a final order reduction. That endpoint is an equation-limit extension: Diffusers refuses a literal zero-terminal non-++ config. EDM uses the ++ algorithms over its existing process.

DPM-Solver++ 2M with no order taper is order=2, algorithm=“dpmsolver++”, solver_type=“midpoint”, lower_order_final=False, euler_at_final=False.

DPMSolverMultistep.init

def init(x, times, process, *, key)

DPMSolverMultistep.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DPMSolverSDE

dataclass source

class DPMSolverSDE(depth: int = MAX_BROWNIAN_DEPTH, seed: int | None = None)

Diffusers 0.34.0’s DPMSolverSDEScheduler, k-diffusion’s sample_dpmpp_sde midpoint solver over a Brownian tree.

Each interval takes two ancestral first-order steps from its own start: one to the geometric midpoint of sigma_t and sigma_s, which the model is read at, and one to sigma_s with that midpoint’s clean prediction. Both steps go down to k-diffusion’s sigma_down and add sigma_up of noise, and both draw that noise from one Brownian path over the trajectory’s sigma interval: the first over [sigma_t, sigma_mid] and the second over [sigma_t, sigma_s], so the two are correlated exactly as nested increments of one path. The source’s sampler transforms sigma with the identity even though its own steps integrate -log(sigma), so the interval widths are sigma differences.

depth resolves the root interval to (sigma_max - sigma_min) / 2**depth: on a published VP table’s span of about 14.6 the default reaches 8.7e-7, at or inside the reference tree’s own 1e-6 tolerance, and it is also where a float32 position runs out of mantissa, so no deeper descent tells two sigmas apart. A zero-sigma target has no ancestral step and lands on the clean prediction.

The root interval is the schedule’s own positive sigma domain, not the extremes of the grid handed to init: the source builds its tree from all the positive sigmas it prepared, so a continuation that walks a suffix of that grid keeps the path the same key gives the whole one. A grid whose only interval lands on sigma zero leaves that domain a single point, which the source also prepares and never queries.

seed is the source’s noise_sampler_seed: with it the tree’s entropy is the checkpoint’s rather than the caller’s, so every walk over the same grid integrates one fixed path however the sampling key changes. It seeds this bridge, not the reference tree, because a Torch seed does not name a JAX stream; what carries over is the contract, a path independent of the walk’s key.

DPMSolverSDE.init

def init(x, times, process, *, key)

DPMSolverSDE.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

DPMSolverSinglestep

dataclass source

class DPMSolverSinglestep(
    order: int = 2,
    algorithm: Literal['dpmsolver++', 'dpmsolver', 'sde-dpmsolver++'] = 'dpmsolver++',
    solver_type: Literal['midpoint', 'heun'] = 'midpoint',
    lower_order_final: bool = False,
)

Diffusers 0.34.0’s grouped DPM-Solver updates from each group’s anchor.

A group of k model evaluations completes one k-th order update. Orders repeat [1, 2] or [1, 2, 3]; an incomplete final group uses lower order, as set_timesteps does in the reference. lower_order_final also lowers the final complete group, and a zero-sigma target forces final order 1. Source-domain checks use that effective order list.

At alpha=0, clean-prediction midpoint and deterministic Heun groups have finite limits. Noise-prediction groups above order 1 and third-order SDE Heun groups diverge; initialization rejects those source/grid pairs.

DPMSolverSinglestep.order_list

def order_list(steps: int) -> list[int]

Reference groups, completing an uneven final group at lower order.

DPMSolverSinglestep.init

def init(x, times, process, *, key)

DPMSolverSinglestep.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

Euler

dataclass source

class Euler()

The DDIM update written as an Euler step of the probability flow ODE. On a variance exploding schedule it is dx/dsigma = eps.

Euler.init

def init(x, times, process, *, key)

Euler.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

EulerAncestral

dataclass source

class EulerAncestral()

Euler with the ancestral noise injection of k-diffusion (get_ancestral_step, eta 1). The step goes down to sigma_down, and sigma_up of fresh noise brings the marginal back to sigma_s. Integrates a GeneralizedNoiseScheduler.

EulerAncestral.init

def init(x, times, process, *, key)

EulerAncestral.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

FlowSDE

dataclass source

class FlowSDE(noise_level: float = 0.7)

Flow-GRPO’s Euler-Maruyama solver on a rectified-flow Process.

Process times may be resolution-shifted. The transition integrates in the resulting physical noise rate, as the reference scheduler does.

FlowSDE.validate

def validate(process: Process) -> None

FlowSDE.init

def init(
    x: jax.Array,
    times: jax.Array,
    process: Process,
    *,
    key: jax.Array,
) -> tuple[()]

FlowSDE.transition

def transition(
    x: jax.Array,
    t: jax.Array,
    t_next: jax.Array,
    denoised: jax.Array,
    eps: jax.Array,
    process: Process,
) -> GaussianTransition

FlowSDE.step

def step(
    x: jax.Array,
    t: jax.Array,
    t_next: jax.Array,
    denoised: jax.Array,
    eps: jax.Array,
    state: tuple[()],
    key: jax.Array,
    process: Process,
    denoise: Callable[[jax.Array, jax.Array], tuple[jax.Array, jax.Array]],
    /,
) -> tuple[jax.Array, tuple[()]]

FlowTrajectory

dataclass source

class FlowTrajectory()

A reverse trajectory, with batch-major states and joint log densities.

states is [batch, points, …], times is [points], and log_probs and stochastic are [batch, points - 1]. Deterministic intervals have NaN log density and a false stochastic mark. The final state is the sample.

GaussianTransition

dataclass source

class GaussianTransition()

An isotropic transition with one variance per batch row.

Sampling and density arithmetic use float32. Densities and KL sum over the sample dimensions. A zero variance is a deterministic transition: sampling returns its mean and log_prob is NaN, since a Dirac measure has no density with respect to Lebesgue measure.

GaussianTransition.sample

def sample(key: jax.Array) -> jax.Array

GaussianTransition.log_prob

def log_prob(value: ArrayLike) -> jax.Array

Joint log density of an observed next state, one value per row.

GaussianTransition.kl

def kl(reference_mean: ArrayLike) -> jax.Array

KL to a reference transition with the same policy-independent variance.

Equal Dirac measures have KL zero; distinct ones have infinite KL.

Generation

dataclass source

class Generation(
    rows: int | None = struct.field(pytree_node=False, default=None),
    decoder: Callable[[ArrayLike, ArrayLike, int], tuple[str, ...]] | None = struct.field(pytree_node=False, default=None),
)

Prompt plus padded continuation, and response-aligned likelihoods.

lengths counts response actions, including EOS. terminated marks a stopping criterion, EOS by default; false marks a length limit. Both log-probability arrays have shape [B, max_new_tokens]. Only positions below lengths are valid. behavior_log_probs describes the distribution that actually drew each action, after the whole transform chain. raw_log_probs describes the unmodified model policy.

A request for n continuations per prompt gives every array [B * n, ...] rows: prompt zero’s n continuations, then prompt one’s. Each row carries its own length, termination and likelihoods.

Arrays keep the placement the task ran with: on a mesh they are global arrays whose rows split over the batch axes, padded to the device count. host() reads this process’s rows real rows back as host arrays. text decodes them through the processor the task was bound to.

text: tuple[str, ...]

Each real row’s valid continuation, decoded on first access.

Generation.host

def host() -> Generation[np.ndarray]

This process’s real rows as host arrays, without the padding a row plan added to fill the devices.

Grammar

dataclass source

class Grammar()

A token automaton on the device.

transitions[state, class] is the state a token of that class leads to, -1 where the token is not allowed; classes[token] is the token’s class. State 0 is the state before the first draw, so a zeroed carry starts a row.

Grammar.start

def start(rows: int) -> jax.Array

Grammar.masked

def masked(state: jax.Array, logits: jax.Array) -> jax.Array

logits [rows, vocab] with the tokens each row’s state forbids at -inf.

Grammar.guiding

def guiding(
    transform: Callable[[StepState, jax.Array], jax.Array],
    state: jax.Array,
) -> Callable[[StepState, jax.Array], jax.Array]

transform behind the mask of each row’s state.

Grammar.advanced

def advanced(state: jax.Array, token: jax.Array, drawn: jax.Array) -> jax.Array

Each row’s state after token; a row that did not draw keeps its state.

A drawn token the state forbids can only come from a transform that forces a token after the mask (ForcedEOS, say), and fails the device check rather than leaving the text outside the language.

Heun

dataclass source

class Heun()

Heun’s second order method (Karras et al. 2022, Algorithm 2): an Euler step, the derivative re-evaluated at its end, and the average of the two.

Diffusers 0.34.0’s HeunDiscreteScheduler limits the clean prediction of both stages under clip_sample; that limit belongs to the process’s conversion, SourceLimitedPrediction, so both evaluations here read the limited prediction without the solver knowing about it.

Heun.init

def init(x, times, process, *, key)

Heun.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

MultiStepDPM

dataclass source

class MultiStepDPM()

A third order multistep integrator of dx/dsigma = eps on a variance exploding schedule, from finite differences of the last three eps.

MultiStepDPM.init

def init(x, times, process, *, key)

MultiStepDPM.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

Sample

dataclass source

class Sample(grammar: Grammar | None = None)

Draw every row independently, one token per step.

This is the loop generate runs when a request names no strategy. Continuations of a prompt share its prefill and run one after another, so decode memory does not grow with n and a routed-expert forward sees the same batch as a single continuation.

grammar holds every draw to a regex or JSON schema (dew.sampling.guided). Each row carries its automaton state through the loop; before a draw the tokens the state forbids score -inf, ahead of the transform chain, so the chain filters and samples inside the language, and the raw likelihood stays the model’s own.

Sampling

dataclass source

class Sampling(
    temperature: float = 1.0,
    top_k: int | None = None,
    eos_id: int | tuple[int, ...] | None = None,
    pad_id: int = 0,
    top_p: float = 1.0,
    min_p: float = 0.0,
)

Token selection and termination. Zero temperature is deterministic argmax.

top_k=None keeps the vocabulary. EOS counts as a sampled action; subsequent output slots contain pad_id and have no likelihood.

A Sampling value compiles to temperature, top-k, top-p and min-p transforms when a request has no explicit logits chain. An explicit chain replaces those transforms. The EOS criterion still joins the request’s stopping criteria.

stops: tuple[int, ...]

The EOS ids that end a draw, none when the policy names no EOS.

Sampling.transforms

def transforms() -> tuple[LogitsTransform, ...]

The complete default chain for a request without explicit transforms.

Zero temperature is the argmax, and the sample-only filters are inactive there, which is what generate() does with do_sample=False.

Sampling.criteria

def criteria() -> tuple[Stopping, ...]

The EOS criterion this policy adds after a caller’s criteria.

Solver

class source

class Solver(Protocol[StateT])

A step of a sampler, and whatever it carries between steps.

StateT is that carried value: nothing for a one-step solver, the previous model outputs for a multi-step one. It is a type parameter, so a solver’s own state type is checked at its call sites.

Solver.init

def init(x, times, process, *, key) -> StateT

Prepare state and check endpoint domains on the concrete time grid.

sample() materializes this grid at compile time, so validation adds no host callbacks to the compiled step. key is the walk’s root key; a solver whose source draws one correlated path over the whole trajectory keeps that path’s state, and every other solver ignores it and draws from the per-step key step is handed.

Every argument is on every solver because this is the surface sample calls. x sizes the carried history (LMS, MultiStepDPM, DEIS, UniPC and the DPM-Solvers), times and process check the grid’s endpoints and count its steps (DDPM, Consistency, DPMSolverSDE, DEIS, UniPC and the DPM-Solvers), and key seeds the Brownian tree of DPMSolverSDE alone. A one-step solver reads none of them and answers ().

Solver.step

def step(
    x,
    t,
    t_next,
    denoised,
    eps,
    state,
    key,
    process,
    denoise,
    /,
) -> tuple[jax.Array, StateT]

x at t_next from x at t and the model’s (denoised, eps) at t. sample passes every argument by position, so a solver over another algebra names the pair for what it reads (the discrete one takes log-probabilities where a Gaussian one takes eps).

Speculative

dataclass source

class Speculative(
    block: int = struct.field(pytree_node=False, default=4),
    confidence: float = struct.field(pytree_node=False, default=0.0),
)

Draft with the model’s prediction depths, verify with the model itself.

The law is algorithm 1 of arXiv 2211.17192, as _speculative_sampling in Transformers 5.16.1 applies it. The first candidate is an ordinary target draw, so it is always accepted, and the model’s prediction depths chain the rest from the target’s last hidden state and each candidate’s embedding, which is what vLLM’s Qwen3_5MultiTokenPredictor does. A proposed x is accepted with probability min(1, p(x) / q(x)) for the target’s post-transform p and the draft’s actual q, compared as a log ratio; the first rejection draws from the normalized positive part of p - q, and a block with nothing rejected draws a bonus token from p. The emitted tokens are therefore distributed exactly as Sample would distribute them, token for token, though not draw for draw at one seed.

Every emitted action, including a replacement or a bonus, records the target’s post-transform log probability as its behaviour and the model’s own log probability as its raw value. The draft’s q, the acceptance probability and the residual are never recorded: none of them is the distribution the emitted action came from.

block candidates per iteration keep every collective the same size. The target cache is saved before the block and the accepted prefix is replayed into it, because a recurrent mixer’s state is a running summary that no cursor can rewind, and the prediction cache is rebuilt the same way. A continuing block emits two or more tokens unless the budget ends first, so ceil(budget / 2) iterations bound the loop.

confidence stops the draft after the first candidate the draft itself is less sure of than that, as the reference’s ConfidenceCriteria does. The later candidates are still computed, at the same shapes, and simply cannot be accepted.

Strategy

class source

class Strategy(Protocol)

The device loop of one generation request.

TextToImage

dataclass source

class TextToImage(
    model: nn.Module,
    process: Process,
    inputs: InputSpec,
    params: Variables,
    autoencoder: AutoEncoder | None = None,
    steps: int = 50,
    guidance: CFG | None = None,
    sampler: Solver[object] = DDIM(),
    grid: Callable[[int], tuple[Process, jax.Array]] | None = None,
    final_denoise: bool = True,
    finish: Callable[[Variables, jax.Array], jax.Array] | None = None,
    blank: Callable[[dict], dict] | None = None,
)

pipe(prompts, seed=0) or pipe(prompts, steps=40, guidance=4.0, sampler=samplers.Heun(), key=key).

params is the objective’s whole tree, the EMA copy merged over the live weights when the run kept one, so a sample comes from the weights a run publishes. steps, guidance and sampler are the defaults a call omits; an objective or a loaded source sets them. grid prepares the process and its explicit time grid for a step count, for a source whose sampler pairs its own sigma and model-time tables; final_denoise False ends a trajectory the way those samplers do. finish runs on the decoded images under the same placement, for a source that ships a checker or an output transform.

Weights keep their placement. On a mesh, prompts split into per-process rows over its batch axes and the result keeps that sharding; each row’s initial noise comes from its global row index, so a pool draws what one process draws for the same prompts.

blank: Callable[[dict], dict] | None

The task’s own unconditional branch in the dtypes of a conditional one, encoded once by whoever built this task (DiffusionObjective.blank_conditions); None encodes it on every call, for a source that has none.

latent_shape: tuple[int, ...]

The per-example shape the model denoises: the sample field’s, or its latent when an autoencoder sits in front of the model.

TextToImage.bind

def bind(variables: Variables) -> TextToImage

Bind another variables snapshot without rebuilding the model or encoders.

TextToImage.from_objective

def from_objective(objective: DiffusionObjective, variables: Variables) -> TextToImage

The objective’s model over variables, sampling the way its evaluation does.

TextToImage.from_run

def from_run(
    directory: str,
    *,
    ema: bool = True,
    step: int | None = None,
    mesh: MeshSpec | None = None,
    layout: Layout | None = None,
    dtype: str | None = None,
    param_dtype: str | None = None,
) -> TextToImage

The run in directory: its run.json built the way the recipe built it, and the weights of its latest checkpoint (or step).

ema reads the averaged weights when the run kept them. With mesh the weights restore straight onto that mesh under layout, the way the trainer places them; without one the default mesh uses the current pool. dtype overrides computation in the model, encoders and VAE. param_dtype overrides parameter storage; None preserves checkpoint storage exactly.

TextToImage.from_pretrained

def from_pretrained(
    repo_id: str,
    *,
    ema: bool = True,
    mesh: MeshSpec | None = None,
    layout: Layout | None = None,
    dtype: str | None = None,
    param_dtype: str | None = None,
) -> TextToImage

A run directory published to the Hugging Face Hub, as dew.interop.hub.push_to_hub(..., raw=True) writes it.

TextToImage.prepared_process

def prepared_process(steps: int) -> tuple[Process, tuple[float, ...] | None]

The process and explicit time grid a steps call walks; the grid is concrete, so the compiled trajectory has its length and values.

TextToImage.prepare

def prepare(
    prompts: str | Sequence[str | Mapping[str, object]],
    *,
    key: jax.Array | None = None,
    seed: int | None = None,
    steps: int | None = None,
    unconditional: str | Sequence[str | Mapping[str, object]] | None = None,
    image: ArrayLike | None = None,
    image_latents: ArrayLike | None = None,
    mask: ArrayLike | None = None,
    noise: ArrayLike | None = None,
    initial: ArrayLike | None = None,
    times: ArrayLike | Sequence[float] | None = None,
    encode_key: jax.Array | None = None,
) -> DenoisingInputs

Encode conditions and construct the initial state on a concrete grid.

Images are uint8 or normalized floating NHWC pixels at the task’s geometry. image_latents skips VAE encoding. A mask adds spatial conditioning to both guidance branches. noise is unit Gaussian noise for noising a clean image; initial is an already-noisy latent state for a continuation or refiner handoff and is never noised again. Explicit times select a partial trajectory in the prepared process. encode_key samples a VAE posterior; None uses its mean.

UniPC

dataclass source

class UniPC(
    order: int = 2,
    solver_type: Literal['bh1', 'bh2'] = 'bh2',
    predict_x0: bool = True,
    lower_order_final: bool = True,
    disable_corrector: tuple[int, ...] = (),
)

UniPC (Zhao et al. 2023, arXiv 2302.04867), Diffusers 0.34.0’s UniPCMultistepScheduler: a unified predictor and corrector in lambda whose weights solve a small linear system over the history’s positions. Each step first corrects the sample the last predictor produced, with the model output just read there and that predictor’s order, then predicts the next sample from the corrected one. solver_type picks B(h) as h (bh1) or e^h - 1 (bh2); predict_x0 integrates the clean prediction, otherwise eps. lower_order_final caps the order by the steps remaining, Diffusers’ rule, so the last step is first order; disable_corrector names the step indices whose predictor’s output is not corrected. The lowered clean-prediction terminal is alpha_t*x_0; epsilon prediction uses the corrected sample with the pre-correction epsilon. Initialization rejects an unlowered higher-order zero target. At an alpha=0 source, bh1 and epsilon correctors diverge unless the first correction is disabled. The finite infinite-node weights use the same Vandermonde system with column scaling.

UniPC.init

def init(x, times, process, *, key)

UniPC.step

def step(x, t, t_next, denoised, eps, state, key, process, denoise)

flow_transition

function source

def flow_transition(
    x: ArrayLike,
    velocity: ArrayLike,
    sigma: ArrayLike,
    sigma_next: ArrayLike,
    *,
    noise_level: float = 0.7,
) -> GaussianTransition

Euler-Maruyama over the physical noise rate, with 0 <= sigma_next <= sigma <= 1.

A rectified flow’s noise rate is its own physical time, which is what FlowSDE reads off the schedule and hands over here. x and velocity are [batch, …]; rates are scalars or [batch]. All density arithmetic is float32. Variance is sigma^2 times the elapsed rate. Invalid rates produce non-finite transitions. At zero noise or zero elapsed rate the result is deterministic.

generate

function source

def generate(
    model: nn.Module,
    params: Variables,
    inputs: ModelInputs | ArrayLike | Sequence[Sequence[int]],
    max_new_tokens: int,
    *,
    key: jax.Array | None = None,
    seed: int | None = None,
    sampling: Sampling = Sampling(),
    n: int = 1,
    logits: Transforms | None = None,
    stopping: Criteria | None = None,
    strategy: Strategy | None = None,
) -> Generation

Generate from numeric model inputs, with an array shorthand for text.

ModelInputs.token_fields[“attention_mask”] identifies real tokens. Missing masks mean all tokens are real. Every row must contain a real token. Prefill evaluates conditioning once; decode reuses the model-owned cache and logical-position state. Each cache compacts real input tokens and leaves paused rows intact.

Parameters keep their placement. On a mesh, rows split over its batch axes and the result keeps that sharding; Generation.host() reads a process’s own rows back. All cooperating processes use the same input shapes, effective decoding components, padding id and continuation count. Decode loops have fixed bounds; any skipped blocks are globally agreed. Keys fold in the global row index and response position, so a pool draws what one process draws for the same rows.

n continuations of each prompt share its prefill and leave as n consecutive rows of every array, in prompt order. Continuation zero of a prompt draws with that prompt’s own key, so n=1 and continuation zero of any larger request are the same draw.

logits is the whole transform chain, in the order it runs. Left as None it is what sampling compiles to, and () runs no transform. stopping adds criteria beside the policy’s EOS one rather than replacing it. strategy replaces the per-row draw loop; None uses Sample.

sample

function source

def sample(
    denoise: Denoiser | DiscreteDenoiser,
    x_T: jax.Array,
    steps: int | None = None,
    *,
    solver: Solver[StateT],
    guidance: CFG | None = None,
    key: jax.Array,
    times: ArrayLike | Sequence[float] | None = None,
    final_denoise: bool = True,
) -> jax.Array

steps points from T to 0: a solver step across each interval, then the model’s clean prediction at the last point.

denoise is process.denoiser(...), which carries the process the solver reads; guidance wraps it. Every step’s noise comes from key folded with the step index, so a trajectory is reproducible from one key. An explicit times grid is the trajectory when given, descending and concrete, for a source whose sampler pairs its own sigma and model-time tables; it decides the length, so a grid of steps + 1 points ending on the terminal is legal and a single point walks nothing. Exactly one of steps and times is passed. final_denoise=False returns the last point’s state without the closing clean prediction, the way those samplers end.

sample_trajectory

function source

def sample_trajectory(
    denoise: Denoiser,
    x_T: jax.Array,
    steps: int,
    *,
    solver: FlowSDE = FlowSDE(),
    guidance: CFG | None = None,
    key: jax.Array,
) -> FlowTrajectory

Record FlowSDE transitions over the same time grid and keys as sample.

steps counts grid points, including both endpoints. A ten-transition rollout therefore uses steps=11. Guidance is applied identically before constructing each Gaussian. Rectified flow’s clean prediction at t=0 is its state, so the last transition already produces the final sample.