dew.inference

Inference tasks: native generation bound to weights, and external engines.

Name	Summary
BlockGeneration	Generates block-diffusion canvases from a DiffusionGemma and its weights. Documented in dew.inference.tasks.
CheckpointBanks	Serve banks by restoring a run’s published weights one bank at a time.
Completion	Holds the choices in prompt-major order, with the SDK responses retained.
DenoisingInputs	Encoded conditioning and initial noise, placed the way a call runs them.
Draw	One sampled continuation of one prompt.
HeldBanks	Serve banks from a variables tree already held in memory.
Images	Decoded samples in [-1, 1], NHWC, keeping the placement the task ran with.
LayerBanks	Serves a banked store’s values, read one bank at a time.
MaskedGeneration	Samples a whole response with native MDLM, holding the prompt fixed. Documented in dew.inference.tasks.
NCCLPush	Publish a policy version to vLLM replicas by NCCL broadcast; see the module docstring.
NativeRolloutServer	Serve rollouts from Dew’s Server on a background stepping thread.
OllamaCompletion	Bind a model to an injected ollama.Client or ollama.AsyncClient.
OpenAICompletion	Bind an OpenAI client, including one configured for a vLLM or SGLang base_url.
OpenAIRolloutServer	Serve rollouts from a vLLM or SGLang OpenAI-compatible completions endpoint.
Processor	Declares the host preprocessing and decoding a loaded source’s processor does. Documented in dew.inference.tasks.
Publication	An engine fleet’s publication as a versioned publisher: load pushes, stamps, then moves version.
RolloutServer	Sample continuations of token prompts under a versioned, reloadable policy.
RunProcessor	Adapts a run’s tokenizer to a task’s host processor.
SafetensorsReload	Publish a policy version to a set of engine replicas through safetensors on disk.
Server	A slot scheduler over a resident KV cache; see the module docstring.
TextGeneration	Generates next tokens from a decoder, its weights and its processor. Documented in dew.inference.tasks.
TextToImage	pipe(prompts, seed=0) or pipe(prompts, steps=40, guidance=4.0, sampler=samplers.Heun(), key=key). Documented in dew.sampling.
Usage	Holds the reported aggregate usage.
VLLMGenerateServer	Serve rollouts from vLLM’s token route, POST /inference/v1/generate.
host_banked	Build the banked store model’s runs read from source’s weights.
pipeline	Load the inference task for source, its weights placed once.

CheckpointBanks

dataclass source

class CheckpointBanks(directory: str, step: int | None = None, ema: bool = False)

Serve banks by restoring a run’s published weights one bank at a time.

ema merges the averaged copy over the live weights, as dew.sampling.pipelines.restore_variables does, so a bank holds what the run publishes. step selects a checkpoint and is resolved to the latest one when it is built, so every bank of one load comes from one checkpoint.

The rows are restored onto the placement the bank asks for with its layer axis dropped and its memory kind set to device, because stacking them is computation and computation reads device memory; the bank the computation writes goes where the store wants it. A load stages one bank’s rows and one bank, and never the model.

stored: Variables

What the checkpoint at step holds, read once, when the source is built.

CheckpointBanks.shapes

def shapes() -> Variables

CheckpointBanks.entry

def entry(placement: Placement) -> Variables

CheckpointBanks.bank

def bank(
    layers: Sequence[int],
    placement: Placement,
    *,
    namespace: tuple[str, ...] = (),
) -> Variables

Completion

dataclass source

class Completion(
    texts: tuple[str, ...],
    finish_reasons: tuple[str | None, ...],
    token_counts: tuple[int | None, ...],
    usage: Usage | None,
    responses: tuple[OllamaResponse | OpenAIResponse, ...],
    tokens: tuple[tuple[int, ...] | None, ...],
    log_probs: tuple[tuple[float, ...] | None, ...],
    routed_experts: tuple[np.ndarray | None, ...] = (),
)

Holds the choices in prompt-major order, with the SDK responses retained.

Per-choice counts stay None when only aggregate usage was reported. finish_reasons retain the backend’s values, including an absent reason. tokens and log_probs hold each choice’s sampled ids and their log-probabilities where the backend reported them: a vLLM completion asked for logprobs with return_tokens_as_token_ids in extra_body reports both. The log-probabilities are whatever distribution the engine was configured to report; this record does not relabel them. routed_experts holds each choice’s [forwarded ids, layers, top_k] expert record when vLLM ran with --enable-return-routed-experts.

DenoisingInputs

dataclass source

class DenoisingInputs(
    conditions: Mapping[str, object] = struct.field(default_factory=dict),
    unconditional: Mapping[str, object] = struct.field(default_factory=dict),
    rows: int | None = struct.field(pytree_node=False, default=None),
    grid_steps: int | None = struct.field(pytree_node=False, default=None),
    process: Process | None = struct.field(pytree_node=False, default=None),
    times: tuple[float, ...] | None = struct.field(pytree_node=False, default=None),
)

Encoded conditioning and initial noise, placed the way a call runs them.

rows counts this process’s real prompts; on a mesh the arrays carry the padded, row-sharded batch a call consumes directly.

Draw

dataclass source

class Draw(
    prompt: tuple[int, ...],
    tokens: tuple[int, ...],
    behavior_log_probs: tuple[float, ...],
    raw_log_probs: tuple[float, ...] | None,
    terminated: bool,
    version: int,
    routed_experts: np.ndarray | None = None,
    support: tuple[tuple[int, ...], ...] | None = None,
)

One sampled continuation of one prompt.

tokens are the sampled actions, EOS included when the draw terminated on it. behavior_log_probs is the likelihood of each action under the distribution that drew it; raw_log_probs is the unmodified model’s, or None when the backend cannot report it. version is the policy version the request was submitted under. routed_experts is the engine’s mixture routing for every id it forwarded and support the ids its sampler kept for each drawn token, when asked for (sessions.Call.routed_experts, sessions.Call.support).

Draw.check_stops

def check_stops(stops: tuple[int, ...]) -> Draw

This draw, refused unless it ends on EOS exactly when it terminated and holds no earlier EOS.

HeldBanks

dataclass source

class HeldBanks(variables: Variables)

Serve banks from a variables tree already held in memory.

The tree is the one model.init or a loader produced, one subtree per layer. Arrays are borrowed: nothing is donated or deleted, so the caller’s tree stays usable. A load therefore holds the source and the growing banked store at once, which costs two copies of the model. Use CheckpointBanks for weights that do not fit twice.

HeldBanks.shapes

def shapes() -> Variables

HeldBanks.entry

def entry(placement: Placement) -> Variables

HeldBanks.bank

def bank(
    layers: Sequence[int],
    placement: Placement,
    *,
    namespace: tuple[str, ...] = (),
) -> Variables

Images

dataclass source

class Images(
    rows: int | None = struct.field(pytree_node=False, default=None),
    latents: ArrayT | None = None,
)

Decoded samples in [-1, 1], NHWC, keeping the placement the task ran with.

host() reads this process’s rows real rows back as a host array.

Images.host

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

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

LayerBanks

class source

class LayerBanks(Protocol)

Serves a banked store’s values, read one bank at a time.

Paths are canonical below each collection. A namespace selects a decoder inside a wrapper; the source knows nothing about how runs are grouped and answers for the layers it is asked for. bank stacks them on a new leading axis in the order given and places the result, so how much memory holding one bank costs is the source’s own business, and is what bounds a load.

LayerBanks.shapes

def shapes() -> Variables

Return the whole store as jax.ShapeDtypeStruct leaves, one per layer.

LayerBanks.entry

def entry(placement: Placement) -> Variables

Read exactly the canonical leaves selected by placement.

LayerBanks.bank

def bank(
    layers: Sequence[int],
    placement: Placement,
    *,
    namespace: tuple[str, ...] = (),
) -> Variables

Return one run’s bank, on those shardings: those layers stacked on a new leading axis in the order given, or, for a run of one layer, that layer’s own subtree. Returned collections are local to namespace.

NCCLPush

dataclass source

class NCCLPush(
    source: Pretrained,
    engines: tuple[str, ...],
    library: str,
    dtype: str = 'bfloat16',
    chunk: int = 512 * 2 ** 20,
    timeout: float = 600.0,
)

Publish a policy version to vLLM replicas by NCCL broadcast; see the module docstring.

source is the Pretrained the replicas were launched from; its export gives the tensors. engines are the replicas’ roots, not their /v1 APIs. library is the engine’s own libnccl.so.2. chunk bounds the bytes of tensors placed on the sender device beside the trainer’s state: the next chunk is copied over while the current one broadcasts, so a push holds at most two chunks there, or two of its largest tensor. Groups open on the first push and stay open until close.

timings holds the seconds each phase of the last push took: gather (the cast and the gather to process 0’s host, every process), and on process 0, which sends, export, send (copies and broadcasts until every replica has loaded the tensors) and total.

NCCLPush.close

def close() -> None

Tear down the groups this side opened; a later push opens them again.

Every group is aborted, not destroyed: NCCL’s destroy finalizes the group, which waits on the engine’s side, and an engine keeps its side open until it exits (the sender waited out jax.distributed’s 300 s shutdown barrier on 4x RTX 3090). A push that failed closes too: a group whose broadcast failed, or whose replica did, cannot carry the next version.

NativeRolloutServer

class source

class NativeRolloutServer(server: Server, *, version: int = 0)

Serve rollouts from Dew’s Server on a background stepping thread.

Submissions and weight loads are serialized with the server’s steps by one lock; a load lands between two steps. The server’s own sampling policy and transforms decide the draws, and a draw keeps both the raw and the behavior likelihood the server records.

NativeRolloutServer.submit

def submit(prompt: Sequence[int], max_new_tokens: int, *, seed: int) -> Future[Draw]

Queue one draw; a request the server refuses raises here and leaves the rest running.

NativeRolloutServer.load

def load(variables: Variables, version: int) -> None

NativeRolloutServer.close

def close() -> None

OllamaCompletion

dataclass source

class OllamaCompletion(model: str, client: OllamaClient | AsyncOllamaClient)

Bind a model to an injected ollama.Client or ollama.AsyncClient.

call/acall accept batches of text and a finite token budget. Additional request fields use the SDK schema, including images, format, context, raw/template/system, think and logprobs; options accepts the SDK Options value or a mapping of backend options. chat/achat support tools and tool results. stream/astream yield native GenerateResponse objects.

OllamaCompletion.stream

def stream(
    prompt: str,
    max_new_tokens: int,
    *,
    seed: int | None = None,
    **parameters: RequestField = {},
) -> Iterator[OllamaResponse]

OllamaCompletion.chat

def chat(
    messages: Sequence[ChatMessage],
    max_new_tokens: int,
    *,
    seed: int | None = None,
    stream: bool = False,
    **parameters: RequestField = {},
) -> OllamaChat | Iterator[OllamaChat]

OllamaCompletion.acall

async def acall(
    prompts: str | Sequence[str],
    max_new_tokens: int,
    *,
    seed: int | None = None,
    **parameters: RequestField = {},
) -> Completion

OllamaCompletion.astream

async def astream(
    prompt: str,
    max_new_tokens: int,
    *,
    seed: int | None = None,
    **parameters: RequestField = {},
) -> AsyncIterator[OllamaResponse]

OllamaCompletion.achat

async def achat(
    messages: Sequence[ChatMessage],
    max_new_tokens: int,
    *,
    seed: int | None = None,
    stream: bool = False,
    **parameters: RequestField = {},
) -> OllamaChat | AsyncIterator[OllamaChat]

OpenAICompletion

dataclass source

class OpenAICompletion(
    model: str,
    client: OpenAI | AsyncOpenAI,
    provider: Literal['openai', 'vllm', 'sglang'] = 'openai',
)

Bind an OpenAI client, including one configured for a vLLM or SGLang base_url.

Native SDK request options pass through to completion/chat resources. Engine-only controls such as top_k/min_p/stop_token_ids, which vLLM and SGLang both accept, belong explicitly in extra_body. SDK responses and streaming chunks retain backend logprobs, token IDs/extensions, tool calls and structured output fields unchanged. Completion prompts are text or token-id rows; a row of ids reaches the engine as ids, with no detokenize/retokenize round trip.

OpenAICompletion.stream

def stream(
    prompts: str | Sequence[str] | TokenRows,
    max_new_tokens: int,
    *,
    seed: int | None = None,
    **parameters: RequestField = {},
) -> Stream[OpenAIResponse]

OpenAICompletion.chat

def chat(
    messages: Sequence[ChatMessage],
    max_new_tokens: int,
    *,
    seed: int | None = None,
    stream: bool = False,
    **parameters: RequestField = {},
) -> ChatCompletion | Stream[ChatCompletionChunk]

OpenAICompletion.acall

async def acall(
    prompts: str | Sequence[str] | TokenRows,
    max_new_tokens: int,
    *,
    seed: int | None = None,
    **parameters: RequestField = {},
) -> Completion

OpenAICompletion.astream

async def astream(
    prompts: str | Sequence[str] | TokenRows,
    max_new_tokens: int,
    *,
    seed: int | None = None,
    **parameters: RequestField = {},
) -> AsyncStream[OpenAIResponse]

OpenAICompletion.achat

async def achat(
    messages: Sequence[ChatMessage],
    max_new_tokens: int,
    *,
    seed: int | None = None,
    stream: bool = False,
    **parameters: RequestField = {},
) -> ChatCompletion | AsyncStream[ChatCompletionChunk]

OpenAIRolloutServer

class source

class OpenAIRolloutServer(
    completion: OpenAICompletion,
    sampling: Sampling,
    weights: WeightSync,
    *,
    version: int = 0,
    workers: int = 64,
    processed_logprobs: bool = False,
    routing: bool = False,
)

Serve rollouts from a vLLM or SGLang OpenAI-compatible completions endpoint.

Each submission is one completion request of token ids, carrying the Sampling policy as engine request controls, a seed, one reported log-probability per sampled token and the ids themselves. workers requests are in flight at once; the engine batches them. The engine is completion.provider.

The reported log-probabilities are the behavior likelihoods only for some policies. vLLM reports raw model log-probabilities unless it runs with --logprobs-mode processed_logprobs, so a transforming policy (temperature other than one) needs processed_logprobs=True to say the engine was started that way. A filtering policy (top-k, top-p or min-p) trains on its recorded support, which only vLLM’s token route returns (VLLMGenerateServer), so vLLM’s completions route refuses one. SGLang’s /v1/completions reports the temperature-scaled distribution before its top-k, top-p and min-p filters and has no field for the filtered one, so this server takes any temperature but no filter. (SGLang’s native /generate reports the filtered likelihood under return_sampling_mask, for a finite top-k; that is the route a filtered policy would need.) SGLANG_RETURN_ORIGINAL_LOGPROB switches the report to raw log-probabilities and must stay unset.

SGLang honors a request’s seed only under --enable-deterministic-inference; otherwise draws are unseeded.

routing=True records vLLM’s routed experts on every draw for routing replay (dew.nn.moe.Routes); the engine runs with --enable-return-routed-experts.

Publication

class source

class Publication(
    weights: WeightSync,
    *,
    version: int = 0,
    stamp: Callable[[int], None] | None = None,
)

An engine fleet’s publication as a versioned publisher: load pushes, stamps, then moves version.

weights is the push (SafetensorsReload, or any WeightSync) and stamp, when set, records a version wherever calls are labelled with one, such as a recording gateway (dew.objectives.rl.harbor.Gateway.stamp). The stamp runs once every replica serves the new version and never before: a gateway stamps each call when its request arrives, so a stamp ahead of any replica would claim weights the call was not sampled from, while one behind them only overstates its lag. A push or stamp that fails raises and leaves version where it was.

Construction stamps version, the version the engines were launched on, so a gateway left at a higher stamp by an earlier run cannot label this run’s first calls with weights it has not served. Every process of a multi-process trainer calls load; process 0 stamps.

Publication.load

def load(variables: Variables, version: int) -> None

RolloutServer

class source

class RolloutServer(Protocol)

Sample continuations of token prompts under a versioned, reloadable policy.

submit never blocks on generation. load replaces the served weights and sets version; submissions after it return draws stamped with it.

RolloutServer.submit

def submit(prompt: Sequence[int], max_new_tokens: int, *, seed: int) -> Future[Draw]

RolloutServer.load

def load(variables: Variables, version: int) -> None

RolloutServer.close

def close() -> None

RunProcessor

dataclass source

class RunProcessor(tokenizer: RunTokenizer)

Adapts a run’s tokenizer to a task’s host processor.

Left-padded prompt rows go in and one string per row comes out.

bos_id: int | None

The id the run’s tokenizer starts a sequence with, or None.

RunProcessor.decode

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

SafetensorsReload

dataclass source

class SafetensorsReload(
    source: Pretrained,
    directory: Path,
    engines: tuple[str, ...],
    engine: Literal['vllm', 'sglang'],
    dtype: str = 'bfloat16',
    timeout: float = 600.0,
)

Publish a policy version to a set of engine replicas through safetensors on disk.

source is the Pretrained the trainer’s model was loaded from; its save writes the weights, the config it derives and the tokenizer files, which is the directory every replica was launched on (a shared filesystem when the replicas are on several hosts). Files are staged beside directory and moved in with os.replace, so no engine reads a half-written file. engines are the replicas’ roots, not their /v1 APIs. An engine call fails the push when it answers other than 200, or when a call that reports its outcome carries success that is not true: both engines report some failures as a 200 with {"success": false}.

One push of version v writes the directory once, then runs the replica sequence on every replica concurrently. vLLM (engine="vllm", checked against v0.30.0), through its development endpoints (VLLM_SERVER_DEV_MODE=1): POST /pause?mode=wait, which lets in-flight requests finish and schedules no new ones; POST /collective_rpc {"method": "reload_weights"}, which reloads from the served directory; POST /reset_prefix_cache, which must answer {"success": true} so no cached prefix outlives the weights that computed it; POST /update_weight_version {"new_version": "v"}, which must answer {"success": true}; and POST /resume. In-flight draws therefore finish wholly on the old weights. A replica that fails after the pause stays paused, so no draw is sampled from weights the push may have half loaded; the next push that succeeds resumes it.

SGLang (engine="sglang", checked against v0.5.20) runs one call per replica, POST /update_weights_from_disk {"model_path": directory, "flush_cache": true, "weight_version": "v"}. SGLang admits it only once every in-flight request has finished, holds new requests until it returns, and flushes the radix cache before answering, so in-flight draws finish wholly on the old weights and no prefix computed by them survives. A load that fails answers 400 with success: false; SGLang’s rollback re-reads the same directory, so the replica then serves whatever that directory holds.

A push with any failed replica raises and names the replicas that did not take v. Publication wraps a push with the version it serves and the stamp a recording gateway needs.

A multi-process trainer calls the push on every process. The pool gathers the served tree to process 0’s host memory (collective_host with held_by="first": every process takes part in the gather, and only process 0 holds the policy), process 0 writes and publishes, and every process learns the outcome at an agreement point, so a failed push raises on all of them instead of leaving the others to hang at the next collective.

SafetensorsReload.write

def write(variables: Variables) -> None

Write variables into directory, file by file atomically; every process of a pool calls it.

Server

class source

class Server(
    model: nn.Module,
    variables: Variables,
    processor: Processor | None,
    *,
    sampling: Sampling,
    transforms: tuple[LogitsTransform, ...],
    stopping: tuple[Stopping, ...],
    grammar: Grammar | None,
    rows: Rows,
    slots: int,
    capacity: int,
    admission: int,
    default_budget: int | None,
    decode_steps: int,
)

A slot scheduler over a resident KV cache; see the module docstring.

Build one with from_task. slots rows run at once over capacity cache slots each; admission bounds the prompts one step prefills. A request whose prompt and budget do not fit the capacity is refused at submit, with the error TextGeneration raises for a request over the model’s context. Over a mesh both counts split evenly over the groups of rows, and so do the pages of a paged pool.

prefix_hits

Prompt tokens served from shared prefix pages instead of prefilled.

cache: Variables

The resident cache, updated in place by every step.

occupancy: int

Rows the host knows to be running, one step behind the device.

Server.from_task

def from_task(
    task: TextGeneration,
    *,
    slots: int,
    capacity: int,
    admission: int | None = None,
    kv_cache: KVCache | None = None,
    chunk: int | None = None,
    prefix_cache: bool = False,
    decode_steps: int = 1,
) -> Server

A server over the task’s model, weights, processor and policy.

capacity rounds up to whole 64-slot tiles, and whole pages of a paged cache, and may not exceed the model’s context. The task’s n has to be one and its strategy the row-wise sampler (with or without a grammar), since the server’s loop is that sampler over rows that come and go.

kv_cache replaces the model’s cache layout (dew.nn.kv_cache). A paged layout pools every row’s pages: pages bounds the memory, and a request is seated once the pool holds its prompt and budget, so short requests fit more rows than pages * page_size / capacity. Over a paged cache, chunk splits a prompt into pieces of at most that many tokens, one piece a step, so a long prompt does not stall the rows decoding beside it, and prefix_cache shares the pages of a prompt prefix an earlier request computed.

Weights placed on a mesh are served on it: the slots, the admission and the pages split over the mesh’s row axes, which have to divide them. Admission defaults to the largest multiple of the group count up to eight rows an iteration. A mesh with a stage or a sequence axis above one, or one over several processes, is refused.

decode_steps runs that many iterations in each device call. The host launches a call’s kernels one after another, which on a tensor axis costs about as long as the step computes, so the devices wait on the host; a call of several iterations pays that once. Requests are seated and draws reach the host at call boundaries: a request waits up to decode_steps iterations for its slot, and a slot a row leaves mid-call stays empty until the next call. The default admission seats decode_steps iterations’ worth of rows a call, so the slots fill as fast. The draws are the same for any value.

Server.reload

def reload(variables: Variables) -> None

Serve variables from the next step on, in place of the current weights.

The tree must match the served one leaf for leaf in shape, and a floating leaf is cast to the served precision (train in float32, serve in bfloat16), so the compiled step runs on unchanged. The leaves are copied onto the served placement: the caller may donate or overwrite its own buffers right after this returns. Rows already running keep their cache and draw their next token from the new weights; a caller that stamps a policy version on a request takes the version it was submitted under. Prompt prefix pages the old weights wrote are no longer shared. Not thread-safe against step: the caller serializes the two.

Server.submit

def submit(
    prompt: Prompt,
    max_new_tokens: int | None = None,
    *,
    key: jax.Array | None = None,
    seed: int | None = None,
) -> Ticket

Queue one request; the ticket resolves to its Generation.

The prompt is validated as TextGeneration validates it, against the server’s capacity in place of the model’s context. A zero budget resolves at once with the prompt alone.

Server.step

def step() -> None

One device call: admit what fits, run decode_steps iterations, read the last call’s.

Server.run

def run() -> None

Step until every queued and running request has resolved.

Usage

dataclass source

class Usage(
    prompt_tokens: int | None = None,
    completion_tokens: int | None = None,
    total_tokens: int | None = None,
)

Holds the reported aggregate usage. None means the backend did not report it.

VLLMGenerateServer

class source

class VLLMGenerateServer(
    base_url: str,
    sampling: Sampling,
    weights: WeightSync,
    *,
    version: int = 0,
    workers: int = 64,
    routing: bool = False,
    timeout: float = 600.0,
)

Serve rollouts from vLLM’s token route, POST /inference/v1/generate.

The one vLLM route that returns, beside the sampled ids and their likelihoods, the ids the sampler kept for each of them (GenerateResponseChoice.sampling_mask, entrypoints/scale_out/token_in_token_out at 1c0eee9), so a top-k or top-p policy trains on its recorded support (support_log_probs). The engine runs with --enable-scale-out (or --tokens-only), --return-sampling-mask (Model Runner V2, no speculative decoding), --logprobs-mode processed_logprobs, so the reported likelihoods are the filtered ones, and --enable-return-routed-experts when routing. vLLM builds the mask only under a finite top-k.

host_banked

function source

def host_banked(
    model: BankedModel,
    source: LayerBanks,
    *,
    mesh: MeshSpec | None = None,
    layout: Layout | None = None,
) -> Variables

Build the banked store model’s runs read from source’s weights.

Each run’s bank is read, stacked and placed on its own, and the copies of one bank are waited for before the next bank is read, so the transfers a load has in flight are one bank’s and not the store’s. Each declared StackView bounds the run length. What the source holds while it answers is the source’s contract, not this one: HeldBanks holds the whole tree it borrowed, CheckpointBanks stages one bank’s rows, and a load of either costs the store plus whatever its source holds. Nothing here donates or deletes a source’s arrays.

layout.host_parameters names the parameters kept in pinned host memory. Only the layers of the stack can be: the stack is what fetches a layer’s parameters as it reaches it, and an embedding table or a head brought over in one piece would cost the device memory it was meant to save, so naming one is refused here, before anything is read. A run whose layers the patterns place differently, leaf for leaf, is refused for the same reason: a bank is one array with one sharding.

pipeline

function source

def pipeline(
    source: str,
    *,
    mesh: MeshSpec | None = None,
    layout: Layout | None = None,
    dtype: str | None = None,
    param_dtype: str | None = None,
    ema: bool = True,
    step: int | None = None,
    revision: str | None = None,
) -> TextToImage | TextGeneration | BlockGeneration | MaskedGeneration

Load the inference task for source, its weights placed once.

source is a run directory, or a source checkpoint directory or Hub repository. mesh places the weights on that mesh under layout (the trainer’s default when None). Without mesh, data parallelism uses the current pool’s devices. dtype selects computation. param_dtype selects parameter storage: None preserves a run’s stored dtypes and uses FP32 masters for a source, and ‘auto’ stores the stored dtypes either way (a source’s config dtype or first floating tensor, as transformers’ dtype=‘auto’ reads it). ema reads a run’s averaged weights; step selects its checkpoint and revision pins a Hub source.

Loading a task also points XLA at the on-disk executable cache, so a restarted process reuses what it already compiled.