dew.interop

Name	Summary
Pretrained	Holds a native model, explicit variables and its checkpoint’s host processor.
Processor	Runs host text and image preprocessing, then normalizes the numeric layout.
dequantize_fp8_blocks	Return float32(weight[i, j]) * float32(scale_inv[i // block, j // block]).
export_run	Write the run in run_dir to destination in its family’s layout.
fp8_format	Return the block size and whether the scales are ue8m0, from quantization.
load_params	Read a safetensors file back into a nested parameter dict.
load_pretrained	Load a source into a native Flax model with explicit parameter trees.
pull_from_hub	Download a snapshot of repo_id and return the directory holding it.
push_to_hub	Upload directory to repo_id, creating the repo when it is missing.
save_hf_layout	Write the weights (save_sharded) and config.json into directory.
save_params	Write a parameter tree to a safetensors file, one tensor per leaf.
save_pretrained_decoder	Write a decoder back out in the HF layout: config.json and its weights, in shards of at most max_shard_size with their index once they exceed it.
split_revision	Split a repo@revision reference into the repo and the revision.
translate_config	Translate one registered family, refusing computation with no counterpart.
translate_weights	Map HF tensors into a CausalTransformer tree.

Pretrained

dataclass source

class Pretrained(
    model: nn.Module,
    variables: Variables,
    processor: Processor | None,
    config: Mapping[str, object],
    source: Path,
    model_config: Mapping[str, object],
    generation_config: Mapping[str, object] = dict(),
    weight_layouts: tuple[WeightLayout, ...] = (),
    retained_tensors: Mapping[str, np.ndarray] = dict(),
    export_adapter: Callable[[nn.Module, Mapping[str, object], Mapping[str, object]], Mapping[str, np.ndarray]] | None = None,
    process: Process | None = None,
    inputs: InputSpec | None = None,
    autoencoder: AutoEncoder | None = None,
    schedule: SourceSchedule | None = None,
    task: SourceTask | None = None,
    finish: Callable[[Mapping[str, object], jax.Array], jax.Array] | None = None,
    quantized_tensors: tuple[str, ...] = (),
    quantized_scale_dtype: str | None = None,
    quantization_grid: Mapping[str, np.ndarray] = dict(),
    revision: str | None = None,
)

Holds a native model, explicit variables and its checkpoint’s host processor.

model_config is the record the model was built from, in Dew’s own vocabulary with the run’s compute dtype and attention kernel, so a caller logs the model it ran.

quantized_scale_dtype: str | None

The dtype a quantized source stored its scales in, where its format leaves that to the checkpoint (DeepSeek-V4’s .scale: float8_e8m0fnu, float32 in the Base releases), so save writes them back in it.

quantization_grid: Mapping[str, np.ndarray]

The scales and zeros an integer format (AWQ, GPTQ) encodes a saved weight against, as the source stored them.

revision: str | None

The Hub commit the source resolved to, whatever branch or tag was asked for; None for a local directory.

layouts: Mapping[str, WeightLayout]

Map each source module name to the layout of its weight: model.<module> for a decoder, unet.<module> for a pipeline component.

These are the names a published adapter file writes, so this is what dew.lora binds a low-rank delta through.

Pretrained.text_generation

def text_generation(
    *,
    sampling: Sampling | None = None,
) -> TextGeneration | MaskedGeneration

Build the text generation task this source describes.

Masked generation refines a full response with Unmask, not the source family’s custom generation recipe. AR sampling overrides are refused.

Without an override the task runs the source’s whole chain, its criteria and the strategy its config names. An explicit sampling replaces the basic policy and clears that chain with it, because the chain was built around the policy the caller just replaced; the criteria and num_return_sequences still come from the source.

Pretrained.block_generation

def block_generation() -> BlockGeneration

Build the DiffusionGemma as a canvas task, defaulting to the source’s sampler config.

Pretrained.text_to_image

def text_to_image() -> TextToImage

Build the latent diffusion source as an image task with its published policy.

Pretrained.export

def export(variables: Mapping[str, object] | None = None) -> Mapping[str, np.ndarray]

The tensors save writes, by their source names; dew.inference.NCCLPush sends these.

A diffusion source writes one set per component, so it has none.

Pretrained.save

def save(
    directory: str | Path,
    *,
    variables: Mapping[str, object] | None = None,
    max_shard_size: int | str = MAX_SHARD_SIZE,
) -> None

Write export’s tensors, in shards of at most max_shard_size, the source’s own config.json, as published, and its tokenizer assets.

Processor

dataclass source

class Processor(
    reference: HostProcessor,
    config: Mapping[str, object],
    record: Mapping[str, object],
    vocab_size: int,
)

Runs host text and image preprocessing, then normalizes the numeric layout.

The checkpoint processor owns resizing, normalization and special-token expansion. Dew organizes its outputs into row-aligned arrays; it does not reproduce the checkpoint’s image preprocessing algorithms.

bos_id: int | None

The id the source’s tokenizer starts a sequence with, or None, read off the tokenizer a processor wraps as _row_padding reads its padding.

Processor.chat

def chat(
    messages: Sequence[Mapping[str, object]],
    *,
    add_generation_prompt: bool = True,
    **template_options: JSON = {},
) -> ModelInputs

Run the source’s actual chat template and processor into numeric inputs.

Template controls such as reasoning_effort and preserve_thinking are interpreted by the checkpoint template. Media-bearing content uses the same reference processor and numeric normalization as plain text.

Processor.from_hf

def from_hf(values: Mapping[str, object]) -> ModelInputs

Validate and normalize actual processor outputs before device use.

Processor.decode

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

Decode token rows with the tokenizer retained by the source processor.

Processor.save_pretrained

def save_pretrained(directory: str | Path) -> None

Save the same processor and tokenizer used by this source.

dequantize_fp8_blocks

function source

def dequantize_fp8_blocks(
    weight: np.ndarray,
    scale_inv: np.ndarray,
    block: int = BLOCK,
) -> np.ndarray

Return float32(weight[i, j]) * float32(scale_inv[i // block, j // block]).

weight may be in any float dtype, scale_inv float32 or E8M0.

export_run

function source

def export_run(
    run_dir: str,
    destination: str | Path,
    *,
    ema: bool = True,
    step: int | None = None,
) -> None

Write the run in run_dir to destination in its family’s layout.

The run loads the way dew.pipeline loads it: run.json for the model record, the latest checkpoint or step for the weights, ema for the averaged copy where the run kept one. The model that comes back decides the layout, and a model with no published layout is refused by name.

The result is a Hugging Face directory: load_pretrained reads it back, and so does transformers for a family it knows.

fp8_format

function source

def fp8_format(quantization: Mapping[str, object]) -> tuple[int, bool]

Return the block size and whether the scales are ue8m0, from quantization.

The finegrained format or refused: E4M3 weights (E4M3_NAMES) in a square block, the scales float32 (scale_fmt absent or ‘float’, V3) or ue8m0 (V3.2). A per-tensor or rectangular scale, or a scale format with no rounding rule here, is not this format.

load_params

function source

def load_params(path) -> ParamTree

Read a safetensors file back into a nested parameter dict.

Leaves are read-only views of the file in their stored dtype, so nothing is placed on a device until the caller asks for it.

load_pretrained

function source

def load_pretrained(
    name_or_dir: str | Path,
    *,
    dtype: str = 'bfloat16',
    param_dtype: str = 'float32',
    attention_impl: str = 'auto',
    max_seq_len: int | None = None,
    revision: str | None = None,
    gguf_file: str | None = None,
    mesh: MeshSpec | None = None,
    layout: Layout | None = None,
    fallback: str | None = None,
) -> Pretrained

Load a source into a native Flax model with explicit parameter trees.

name_or_dir is a local HF directory or a Hub model identifier. The decoder/tower/projector maps preserve their established internal paths; wrapper variables join under their existing component names. Processor artifacts are loaded only when the source contains them. dtype selects computation; param_dtype independently selects floating parameter storage and defaults to FP32 masters, or ‘auto’ stores the checkpoint’s own dtype (_checkpoint_dtype). Frozen component weights (text encoders and VAE) follow it too; router, clipping, positional and safety state retain their own FP32/integer contracts.

Without mesh or layout the variables are host arrays. With either, they are placed on that mesh (the default MeshSpec() when only layout is given) under that layout, one leaf at a time: a decoder’s leaves are read from the mapped checkpoint one device shard at a time and cast and transposed there (dew.interop.streaming), so the host never holds the translated model. Towers, projectors and a quantized source’s dequantized tensors are still built whole on the host first.

gguf_file names a GGUF file in the repo or directory: its metadata is the config, its block-quantized tensors are dequantized to float32 (dew.interop.gguf), and its tokenizer is the processor where the repo ships no tokenizer.

fallback="torchax" opts into tier 3 for any causal LM transformers can build, registered or not: transformers’ PyTorch forward lowered to JAX by torchax (dew.interop.torchax_fallback), with no Dew kernels, sharding rules or cached generation.

pull_from_hub

function source

def pull_from_hub(repo_id: str, revision: str | None = None) -> Path

Download a snapshot of repo_id and return the directory holding it.

revision is a branch, tag or commit; None takes the default branch. The path is inside the hub cache, so a second call with the same revision downloads nothing.

push_to_hub

function source

def push_to_hub(
    directory,
    repo_id: str,
    *,
    private: bool = False,
    commit_message: str = 'Upload dew export',
    raw: bool = False,
) -> None

Upload directory to repo_id, creating the repo when it is missing.

The files land at the root of the repo under the names they have on disk, so an export written by save_hf_layout arrives as the model.safetensors and config.json pair a Hugging Face loader looks for.

A run directory is exported first, because the orbax checkpoint and the run.json beside it are Dew’s own format and nothing on the Hub reads them: what goes up is what export_run writes. raw uploads the run directory itself instead, which is the form TextToImage.from_run and the text tasks’ from_pretrained pull back.

save_hf_layout

function source

def save_hf_layout(
    params,
    config: Mapping[str, object],
    directory,
    max_shard_size: int | str = MAX_SHARD_SIZE,
) -> None

Write the weights (save_sharded) and config.json into directory.

That is what a Hugging Face style loader looks for. params is a flat table of named tensors, or a tree whose ’/‘-joined paths name them. The config is written as given. Dew does not translate its own config vocabulary into anyone else’s.

save_params

function source

def save_params(params, path) -> None

Write a parameter tree to a safetensors file, one tensor per leaf.

save_pretrained_decoder

function source

def save_pretrained_decoder(
    model,
    variables,
    directory,
    *,
    tokenizer: str | ExportTokenizer | None = None,
    generation_config: Mapping[str, object] | None = None,
    max_shard_size: int | str = MAX_SHARD_SIZE,
) -> None

Write a decoder back out in the HF layout: config.json and its weights, in shards of at most max_shard_size with their index once they exceed it.

Derive the config from native computation and encode all variable collections through the matching family. Source-bound exports instead retain their source layout in Pretrained.save. Gemma4 writes frozen or trainable layer-scalar values into HF buffers; reloading that layout preserves computation, not the native scalar training policy.

tokenizer is the vocabulary the weights were trained against, by object or by name; save_export_assets writes its files beside them, so one call leaves a directory load_pretrained reads back with its processor. Pretrained.save writes a decoder’s weights through the same encoder (Pretrained.export), so the two leave the same weights behind.

split_revision

function source

def split_revision(source: str) -> tuple[str, str | None]

Split a repo@revision reference into the repo and the revision.

Levanter’s RepoRef spelling: a branch, tag or commit after the last ’@’. A local directory, or a reference without ’@’, names no revision. Hub repo ids cannot contain ’@’.

translate_config

function source

def translate_config(hf_config: Mapping[str, object]) -> DecoderFields

Translate one registered family, refusing computation with no counterpart.

translate_weights

function source

def translate_weights(
    hf_tensors: Mapping[str, np.ndarray],
    config: DecoderFields,
    model_type: str | None = None,
    *,
    param_dtype: str = 'float32',
    lazy: bool = False,
) -> Variables

Map HF tensors into a CausalTransformer tree. Parameters default to FP32.

Linear weights arrive as [out, in] and nn.Dense keeps [in, out], so every .kernel is transposed; norm .weight becomes .scale; Gemma’s post_attention_layernorm and post_feedforward_layernorm land on the sandwich norms, where Gemma applies them.

A tied checkpoint carries lm_head.weight as well, as a copy of the embedding (Qwen3-0.6B does). The copy is checked and dropped. The tree has one leaf for the two, and a checkpoint whose “tied” head is a different matrix would otherwise load as a model that computes something else. DeepSeek’s routed experts arrive one tensor per expert and stack onto an expert dimension here; its dense shared experts, MLA projections and indexer map by pattern like everything else, and its routers’ balancing bias lands in the moe collection beside params. param_dtype changes floating parameter storage, independently of compute dtype. Router and frozen state remain FP32; integer indices retain their native dtype. Conversion happens per leaf before its layout copy.

With lazy every leaf is a SourceLeaf over the stored tensors, read only when it is placed (dew.interop.streaming); otherwise each is read whole here.

model_type names the source’s own family where the caller read it off a config.json. Without it the family comes from the record, which is what the backbone would be built from and so cannot tell two families apart that compute the same thing under different tensor names: Kimi K2.5’s decoder is DeepSeek V3’s computation nested under language_model..