ashishkumar4/dew

How ashishkumar4/dew works

Dew is a JAX/Flax framework for training language, diffusion, and JEPA models, with shared data, optimization, sharding, evaluation, and checkpoint facilities. The main workflow is to define data, a model and objective, then run Trainer and sample or otherwise use the resulting state. The README explicitly demonstrates the diffusion path; source-index references and excerpts support the shared data and interoperability links. Several listed capabilities are represented as subsystems, but their internal wiring is not inferred where source evidence was not shown.

Main entry points: __init__.py Trainer · pipeline.py Record pipelines · compile.py Pipeline compiler · causal_transformer.py Model backbones

Components 21

Data preparation

Data sources src/dew/data/sources/base.py
Record pipelines src/dew/data/pipeline.py
Text processing src/dew/data/text.py
Pipeline compiler src/dew/data/compile.py
Data loader src/dew/data/dataset.py

Training core

Batch contract src/dew/batch.py
Trainer src/dew/training/__init__.py
Training state src/dew/state.py
Checkpoints src/dew/checkpoints/__init__.py

Models and objectives

Model backbones src/dew/nn/backbones/causal_transformer.py
Training objectives src/dew/objectives/base.py
Diffusion objective src/dew/objectives/diffusion/objective.py
Sampling src/dew/sampling/__init__.py
Pipeline assembly src/dew/interop/pipeline_assembly.py
Decoder families src/dew/interop/decoder_parts.py

Runtime capabilities

Inference serving src/dew/inference/serving.py
Evaluation harness src/dew/eval/harness.py
Device placement src/dew/distributed/placement.py
Kernel algebra src/dew/algebra/program.py
Model registry src/dew/registry.py

Other components

Framework user

Connections 19

Framework user → Data sources: configures data
Framework user → Model backbones: chooses model
Framework user → Diffusion objective: configures objective
Data sources → Record pipelines: feeds records src/dew/data/pipeline.py
Record pipelines → Text processing: adds text steps src/dew/data/presets.py
Record pipelines → Pipeline compiler: lowers to src/dew/data/pipeline.py
Pipeline compiler → Data loader: builds loader src/dew/data/compile.py
Pipeline compiler → Batch contract: forms batches src/dew/data/compile.py
Model backbones → Batch contract: reads batches src/dew/batch.py
Training objectives → Batch contract: scores batches src/dew/batch.py
Diffusion objective → Model backbones: uses model
Framework user → Trainer: runs training
Trainer → Data loader: fits data
Trainer → Checkpoints: writes checkpoints src/dew/training/__init__.py
Trainer → Training state: returns state
Training state → Sampling: supplies weights
Sampling → Diffusion objective: uses process
Pipeline assembly → Decoder families: builds components src/dew/interop/pipeline_assembly.py
Decoder families → Pipeline assembly: supports assembly src/dew/interop/decoder_parts.py

Drawn by GitDiagram on November 13, 2026 from the repository's file tree, README and source files. Read as Markdown ·