dew.inference.tasks

Bind a model, its weights and its host processing into a reusable task.

A task binds what a generation needs beyond the request itself: the native model, a captured variables mapping and, when the source ships one, the host processor that turns text and media into ModelInputs. Controls stay the typed values training already uses (Sampling, BlockProcess); results stay the typed records the kernels produce. bind gives the same task over other weights, which is how a training loop draws from a policy snapshot without holding the trainer’s mutable mapping. Array buffers are shared; do not mutate, donate or delete those buffers while the task is using them.

Name	Summary
Rows
Request
SHAPE_BUCKETS	The shapes a text request is rounded up to: powers of two.
Processor	Declares the host preprocessing and decoding a loaded source’s processor does.
run_record	Read the run.json a run directory publishes beside its checkpoints.
TextGeneration	Generates next tokens from a decoder, its weights and its processor.
BlockGeneration	Generates block-diffusion canvases from a DiffusionGemma and its weights.
MaskedGeneration	Samples a whole response with native MDLM, holding the prompt fixed.

Rows

attribute source

Rows = ModelInputs | ArrayLike | Sequence[Sequence[int]]

Request

attribute source

Request = str | Sequence[str] | Rows

SHAPE_BUCKETS

attribute source

SHAPE_BUCKETS = tuple(1 << exponent for exponent in range(21))

The shapes a text request is rounded up to: powers of two.

Every distinct prompt width, batch, budget and continuation count is its own several-second XLA compile, and served requests are rarely the same length twice. So a prompt pads left to the smallest bucket of 64 or more, where the attention mask hides the filler as it already hides the padding a ragged batch needs; a budget rounds up to the next power of two, so a one-token request scans one trip and a scoring probe pays no decode it did not ask for, and the trips past the request come off the result; the cache for the call is the two together, rounded up again, in place of the model’s whole max_seq_len. Rows are the caller’s and are not bucketed. A request whose buckets would need more capacity than the model’s max_seq_len keeps its own shapes, so the ceiling refuses what it refuses today.

Processor

class source

class Processor(Protocol)

Declares the host preprocessing and decoding a loaded source’s processor does.

Processor.decode

def decode(tokens: ArrayLike) -> list[str]

run_record

function source

def run_record(directory: str) -> Mapping[str, object]

Read the run.json a run directory publishes beside its checkpoints.

TextGeneration

dataclass source

class TextGeneration(
    model: nn.Module,
    variables: Variables,
    processor: Processor | None = None,
    sampling: Sampling = Sampling(),
    max_new_tokens: int | None = None,
    max_length: int | None = None,
    n: int = 1,
    logits: tuple[LogitsTransform, ...] | None = None,
    stopping: tuple[Stopping, ...] = (),
    strategy: Strategy | None = None,
)

Generates next tokens from a decoder, its weights and its processor.

A call runs the shared cached prefill and decode; the result is the Generation a training rollout consumes, with the actual and raw-policy likelihood of every drawn action. sampling is the policy a call uses when it passes none, max_new_tokens the budget and n the number of continuations per prompt; a loaded source fills all three from its generation config. n continuations of a prompt leave as n consecutive rows, in prompt order. Weights keep their placement: on a mesh, rows split over its batch axes and results keep that sharding.

logits is the whole transform chain, stopping the criteria that run beside the policy’s EOS one, and strategy the device loop. logits=None means the chain sampling compiles to. A call replaces each of them whole, so a caller that wants to add to a bound chain writes logits=task.logits + (mine,), and an explicit sampling= on a call replaces a bound chain with its own, because the policy it overrides is what that chain was built from.

A call runs at SHAPE_BUCKETS shapes over a cache the bucket sizes, and hands back the shapes the request asked for, so two requests of nearby lengths share one compiled executable and neither pays for the model’s whole context.

TextGeneration.bind

def bind(variables: Variables) -> TextGeneration

Return the same task over other weights, such as a policy snapshot.

TextGeneration.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,
) -> TextGeneration

Load the causal run in directory: the model its run.json records, rebuilt the way the recipe built it, over the weights of its latest checkpoint (or step), decoding through the run’s own tokenizer.

ema reads the averaged weights, except under an objective whose average is a reference policy rather than the trained one. With mesh the weights restore straight onto that mesh under layout, the way the trainer places them. dtype overrides computation; param_dtype overrides parameter storage, and None preserves what the checkpoint stored. The run’s preview budget and sampling policy become the task’s defaults.

TextGeneration.from_pretrained

def from_pretrained(
    repo_id: 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,
) -> TextGeneration

Load a run directory published to the Hugging Face Hub.

dew.interop.hub.push_to_hub(..., raw=True) is what writes it.

TextGeneration.decode

def decode(generation: Generation) -> tuple[str, ...]

Return each row’s valid continuation as text, empty without a processor.

BlockGeneration

dataclass source

class BlockGeneration(
    model: DiffusionGemma,
    variables: Variables,
    process: BlockProcess,
    processor: Processor | None = None,
    eos_token_ids: tuple[int, ...] = (),
    pad_token_id: int = 0,
    max_new_tokens: int | None = None,
    max_length: int | None = None,
    n: int = 1,
)

Generates block-diffusion canvases from a DiffusionGemma and its weights.

A call runs prefill, refinement and clean-token commits as one device computation; the CanvasGeneration result carries no autoregressive likelihoods. process is the published sampler configuration used when a call passes none, max_new_tokens the budget a call omits and n the number of continuations per prompt, which leave as n consecutive rows in prompt order.

BlockGeneration.bind

def bind(variables: Variables) -> BlockGeneration

Return the same task over other weights.

BlockGeneration.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,
) -> BlockGeneration

Load the block-diffusion run in directory: the DiffusionGemma its run.json records over the weights of its latest checkpoint (or step), sampling over the canvas the model declares.

The arguments carry what TextGeneration.from_run carries: ema selects the averaged weights, mesh and layout place them, and the two dtypes override computation and storage.

BlockGeneration.from_pretrained

def from_pretrained(
    repo_id: 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,
) -> BlockGeneration

Load a run directory published to the Hugging Face Hub.

dew.interop.hub.push_to_hub(..., raw=True) is what writes it.

BlockGeneration.decode

def decode(generation: CanvasGeneration) -> tuple[str, ...]

Return each row’s valid continuation as text, empty without a processor.

MaskedGeneration

dataclass source

class MaskedGeneration(
    model: nn.Module,
    variables: Variables,
    process: DiscreteProcess,
    processor: Processor | None = None,
    sampler: Unmask = Unmask(),
    steps: int = MDLM_STEPS,
    eos_token_ids: tuple[int, ...] = (),
    pad_token_id: int = 0,
    max_new_tokens: int | None = None,
    max_length: int | None = None,
    n: int = 1,
)

Samples a whole response with native MDLM, holding the prompt fixed.

This is not LLaDA’s or Dream’s source-specific remasking recipe. EOS trims the finished response, not the bidirectional denoising trajectory. Results carry refinement counts, never autoregressive action likelihoods.

MaskedGeneration.bind

def bind(variables: Variables) -> MaskedGeneration

Return the same native MDLM task over another weight snapshot.

MaskedGeneration.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,
) -> MaskedGeneration

Load the masked-diffusion run in directory: the bidirectional model its run.json records over the weights of its latest checkpoint (or step), refined with MDLM over the run’s own mask token.

The arguments carry what TextGeneration.from_run carries, and the run’s preview budget becomes the response length a call omits.

MaskedGeneration.from_pretrained

def from_pretrained(
    repo_id: 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,
) -> MaskedGeneration

Load a run directory published to the Hugging Face Hub.

dew.interop.hub.push_to_hub(..., raw=True) is what writes it.

MaskedGeneration.decode

def decode(generation: CanvasGeneration) -> tuple[str, ...]

Return each row’s valid response as text, empty without a processor.