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. Documented in dew.sampling.solvers.
DDPM	Exact ancestral sampler for the reverse diffusion SDE. Documented in dew.sampling.solvers.
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. Documented in dew.sampling.solvers.
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. Documented in dew.sampling.solvers.
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. Documented in dew.sampling.solvers.
PNDM	PNDM (Liu et al. Documented in dew.sampling.solvers.
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. Documented in dew.sampling.solvers.
TCD	Trajectory consistency sampling (Zheng et al. Documented in dew.sampling.solvers.
Beam	Deterministic beam search over one shared prefill.
Consistency	Multistep consistency sampling (Song et al. Documented in dew.sampling.solvers.
DPMSolverMultistep	DPM-Solver (Lu et al. Documented in dew.sampling.solvers.
DPMSolverSDE	Diffusers 0.34.0’s DPMSolverSDEScheduler, k-diffusion’s sample_dpmpp_sde midpoint solver over a Brownian tree. Documented in dew.sampling.solvers.
DPMSolverSinglestep	Diffusers 0.34.0’s grouped DPM-Solver updates from each group’s anchor. Documented in dew.sampling.solvers.
Euler	The DDIM update written as an Euler step of the probability flow ODE. Documented in dew.sampling.solvers.
EulerAncestral	Euler with the ancestral noise injection of k-diffusion (get_ancestral_step, eta 1). Documented in dew.sampling.solvers.
FlowSDE	Flow-GRPO’s Euler-Maruyama solver on a rectified-flow Process. Documented in dew.sampling.flow.
FlowTrajectory	A reverse trajectory, with batch-major states and joint log densities. Documented in dew.sampling.flow.
GaussianTransition	An isotropic transition with one variance per batch row. Documented in dew.sampling.flow.
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. Documented in dew.sampling.solvers.
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. Documented in dew.sampling.solvers.
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. Documented in dew.sampling.solvers.
Speculative	Draft with the model’s prediction depths or its block drafter, 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. Documented in dew.sampling.solvers.
flow_transition	Euler-Maruyama over the physical noise rate, with 0 <= sigma_next <= sigma <= 1. Documented in dew.sampling.flow.
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. Documented in dew.sampling.flow.

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.

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.

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.

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.

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 or its block drafter, 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 block drafter (DeepSeek-V4.1’s DSpark) drafts them in one pass after the first, each drawn from its position’s logits as the pass reaches it. 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.

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.