dew.registry

Name the things a run is made of.

One Registry per kind, including model components such as mixers, towers and projectors. A registry is a decorator, a mapping and an attribute view over the same table, so models["simple_dit"], models.SimpleDiT and the class are one object. A name or a field the table does not know raises.

The registries are empty at import. Each member registers itself where it is defined, so importing a package fills its table and the registry module imports none of them.

Name	Summary
`REGISTRIES`
`Registry`	Names one kind of thing: a decorator, a mapping and an attribute view.
`datasets`
`dtype_name`	Return the name `resolve_dtype` accepts for a dtype, for a logged config.
`encoders`
`metrics`
`mixers`
`models`
`objectives`
`presets`
`projectors`
`resolve_dtype`	Resolve a dtype for a module field: a jnp dtype, one of its names, or None.
`samplers`
`schedules`
`towers`
`with_precision`	Return a model config with the run’s compute dtype and attention kernel in it.

REGISTRIES

attribute source

REGISTRIES = (models, presets, samplers, datasets, encoders, metrics, objectives, mixers, towers, projectors, schedules)

Registry

class source

class Registry(kind: str, *, record: Literal['name', 'kind'] = 'name')

Names one kind of thing: a decorator, a mapping and an attribute view.

union: type[Built] | types.UnionType: Return Union[...] of the members, for a tyro subcommand over the table.

Registry.name_of

def name_of(member: Named) -> str

Return the name a member was registered under. The table is scanned by identity, so this takes a member of any registry, whatever it makes.

Registry.build

def build(
    name: str,
    record: Mapping[str, object] = NO_RECORD,
    /,
    **fields: Configured = {},
) -> Built

Construct the member called name from a record, keyword fields, or both.

A field the member does not declare is an error. Fields arrive from JSON as often as from code, so a field whose declared type is a value builds from a record here, where a logged config becomes an object: models.build("m", attention={"heads": 8}) and models.build("m", attention=Attention(heads=8)) agree.

A whole parsed config is the positional record: its values are unnarrowed, and narrowing them against the member’s declared types is this method’s job, so the splat happens here rather than at a caller that would have to know the member’s fields to write it.

datasets

attribute source

datasets: Registry[type[DatasetSpec], DatasetSpec] = Registry('dataset')

Registered name	Object
`array_record_images`	`ArrayRecordImages`
`cc12m`	`CC12M`
`cc3m`	`CC3M`
`chat_messages`	`ChatMessages`
`combined_30m`	`Combined30M`
`combined_msml612`	`CombinedMsml612`
`combined_online`	`CombinedOnline`
`diffusiondb`	`DiffusionDB`
`hf`	`HubDataset`
`hf_images`	`HFImages`
`laion12m_coco`	`Laion12mCoco`
`laion2b_aesthetic`	`Laion2bAesthetic`
`local_videos`	`LocalVideos`
`oxford_flowers102`	`OxfordFlowers`
`packed_tokens`	`PackedTokens`
`preference_pairs`	`PreferencePairs`
`prompts`	`Prompts`
`tfds`	`PreparedTFDS`
`token_windows`	`TokenWindows`
`voxceleb2`	`VoxCeleb2`

dtype_name

function source

def dtype_name(value: DTypeLike | None) -> DtypeName | None

Return the name resolve_dtype accepts for a dtype, for a logged config.

encoders

attribute source

encoders: Registry[type[ConditionEncoder[Any]], ConditionEncoder[Any]] = Registry('encoder')

Registered name	Object
`char_table`	`CharTable`
`clip_text`	`CLIPText`
`diffusion_text`	`DiffusionConditioner`
`qwen_image_text`	`dew.inputs.diffusion.QwenImageConditioner`
`t5`	`T5Text`

metrics

attribute source

metrics: Registry[Callable[..., Metric], Metric] = Registry('metric')

Registered name	Object
`clip`	`clip`
`clip_score`	`clip_score_metric`
`fid`	`FID`
`knn_probe`	`knn_probe`
`linear_probe`	`linear_probe`
`perplexity`	`perplexity`
`psnr`	`psnr`
`ssim`	`ssim`

mixers

attribute source

mixers: Registry[type[MixerBase], MixerBase] = Registry('mixer', record='kind')

Registered name	Object
`attention`	`dew.nn.mixers.attention.AttentionMixer`
`deepseek_v4`	`dew.nn.deepseek_v4.DeepseekV4Mixer`
`gated_delta_net`	`dew.nn.mixers.gated_delta_net.GatedDeltaNetMixer`
`kimi_delta_attention`	`dew.nn.kda.KimiDeltaAttentionMixer`
`kpool_sparse_attention`	`dew.nn.dsa_kpool.KPoolSparseAttentionMixer`
`llama4`	`dew.nn.llama4.Llama4Mixer`
`mamba2`	`dew.nn.mixers.mamba2.Mamba2Mixer`
`mla`	`dew.nn.mla.MLAMixer`

models

attribute source

models: Registry[type[nn.Module], nn.Module] = Registry('model')

Registered name	Object
`causal_transformer`	`CausalTransformer`
`diffusion_gemma`	`DiffusionGemma`
`flux_transformer`	`FluxTransformer`
`hierarchical_mmdit`	`HierarchicalMMDiT`
`hybrid_dit`	`HybridSSMAttentionDiT`
`jepa_encoder`	`JepaEncoder`
`jepa_predictor`	`JepaPredictor`
`jepa_video_encoder`	`JepaVideoEncoder`
`multimodal_transformer`	`dew.nn.multimodal.MultimodalTransformer`
`qwen_image_transformer`	`QwenImageTransformer`
`sd3_transformer`	`SD3Transformer`
`simple_dit`	`SimpleDiT`
`simple_mmdit`	`SimpleMMDiT`
`simple_udit`	`SimpleUDiT`
`unet`	`Unet`
`unet_2d_condition`	`UNet2DCondition`
`unet_3d`	`UNet3D`
`uvit`	`UViT`
`video_dit`	`VideoDiT`

objectives

attribute source

objectives: Registry[type[Objective], Objective] = Registry('objective')

Registered name	Object
`block_diffusion`	`BlockDiffusionObjective`
`diffusion`	`DiffusionObjective`
`distillation`	`DistillationObjective`
`dpo`	`DPOObjective`
`flow_grpo`	`FlowGRPOObjective`
`grpo`	`GRPOObjective`
`jepa`	`JepaObjective`
`lm`	`LMObjective`
`masked_diffusion`	`MaskedDiffusionObjective`
`ppo`	`PPOObjective`

presets

attribute source

presets: Registry[type[Preset], Preset] = Registry('preset')

Registered name	Object
`cosine`	`Cosine`
`edm`	`EDM`
`flow`	`Flow`
`karras`	`Karras`
`mdlm`	`MDLM`
`sqrt`	`Sqrt`

projectors

attribute source

projectors: Registry[type[ProjectorBase], ProjectorBase] = Registry('projector', record='kind')

Registered name	Object
`gemma`	`dew.nn.vision.GemmaProjector`
`gemma3n`	`dew.nn.vision.Gemma3nProjector`
`gemma4`	`dew.nn.vision.Gemma4Projector`
`llama4`	`dew.nn.vision.Llama4Projector`
`qwen3_5`	`dew.nn.vision.Qwen35Projector`

resolve_dtype

function source

def resolve_dtype(value: object) -> DTypeLike | None

Resolve a dtype for a module field: a jnp dtype, one of its names, or None.

Every field named dtype is read here, wherever it arrives from, so a dtype passes through and a name becomes the dtype it names. Anything else is refused here rather than inside a module’s first cast.

samplers

attribute source

samplers: Registry[type[Solver[Any]], Solver[Any]] = Registry('sampler')

Registered name	Object
`consistency`	`Consistency`
`ddim`	`DDIM`
`ddpm`	`DDPM`
`deis`	`DEIS`
`dpmsolver_multistep`	`DPMSolverMultistep`
`dpmsolver_sde`	`DPMSolverSDE`
`dpmsolver_singlestep`	`DPMSolverSinglestep`
`euler`	`Euler`
`euler_ancestral`	`EulerAncestral`
`flow_sde`	`FlowSDE`
`heun`	`Heun`
`kdpm2`	`KDPM2`
`lms`	`LMS`
`multistep_dpm`	`MultiStepDPM`
`pndm`	`PNDM`
`rk4`	`RK4`
`tcd`	`TCD`
`unipc`	`UniPC`
`unmask`	`Unmask`

schedules

attribute source

schedules: Registry[type[ScheduleBase], ScheduleBase] = Registry('schedule', record='kind')

Registered name	Object
`cosine`	`Cosine`
`linear`	`Linear`
`power`	`Power`

towers

attribute source

towers: Registry[type[TowerBase], TowerBase] = Registry('tower', record='kind')

Registered name	Object
`gemma3n`	`dew.nn.vision.Gemma3nVision`
`gemma3n_audio`	`dew.nn.audio.Gemma3nAudio`
`gemma4`	`dew.nn.vision.Gemma4Vision`
`gemma4_audio`	`dew.nn.audio.Gemma4Audio`
`llama4`	`dew.nn.vision.Llama4Vision`
`qwen3_5`	`dew.nn.vision.Qwen35Vision`
`siglip`	`dew.nn.vision.SiglipVision`

with_precision

function source

def with_precision(
    name: str,
    config: Mapping[str, object],
    *,
    dtype: str,
    attention_impl: str,
    param_dtype: str | None = None,
    matmul_precision: str | None = None,
) -> Mapping[str, object]

Return a model config with the run’s compute dtype and attention kernel in it.