# ashishkumar4/dew architecture

> Architecture diagram of the GitHub repository ashishkumar4/dew, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-11-13). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/ashishkumar4/dew
- Repository: https://github.com/ashishkumar4/dew

## Overview

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: [`src/dew/training/__init__.py`](https://github.com/ashishkumar4/dew/blob/main/src/dew/training/__init__.py) (Trainer), [`src/dew/data/pipeline.py`](https://github.com/ashishkumar4/dew/blob/main/src/dew/data/pipeline.py) (Record pipelines), [`src/dew/data/compile.py`](https://github.com/ashishkumar4/dew/blob/main/src/dew/data/compile.py) (Pipeline compiler), [`src/dew/nn/backbones/causal_transformer.py`](https://github.com/ashishkumar4/dew/blob/main/src/dew/nn/backbones/causal_transformer.py) (Model backbones)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data preparation"]
  node_sources["Data sources<br/>[base.py]"]
  node_pipeline["Record pipelines<br/>[pipeline.py]"]
  node_text["Text processing<br/>[text.py]"]
  node_compiler["Pipeline compiler<br/>[compile.py]"]
  node_loader["Data loader<br/>[dataset.py]"]
end

subgraph group_train["Training core"]
  node_batch["Batch contract<br/>[batch.py]"]
  node_trainer["Trainer<br/>[__init__.py]"]
  node_state["Training state<br/>[state.py]"]
  node_checkpoints["Checkpoints<br/>[__init__.py]"]
end

subgraph group_models["Models and objectives"]
  node_backbones["Model backbones"]
  node_objectives["Training objectives<br/>[base.py]"]
  node_diffusion["Diffusion objective<br/>[objective.py]"]
  node_sampling["Sampling<br/>[__init__.py]"]
  node_interop["Pipeline assembly"]
  node_decoder["Decoder families<br/>[decoder_parts.py]"]
end

subgraph group_runtime["Runtime capabilities"]
  node_serving["Inference serving<br/>[serving.py]"]
  node_evaluation["Evaluation harness<br/>[harness.py]"]
  node_distributed["Device placement<br/>[placement.py]"]
  node_algebra["Kernel algebra<br/>[program.py]"]
  node_registry["Model registry<br/>[registry.py]"]
end

node_user(("Framework user"))

node_user -->|"configures data"| node_sources
node_user -->|"chooses model"| node_backbones
node_user -->|"configures objective"| node_diffusion
node_sources -->|"feeds records"| node_pipeline
node_pipeline -->|"adds text steps"| node_text
node_pipeline -->|"lowers to"| node_compiler
node_compiler -->|"builds loader"| node_loader
node_compiler -->|"forms batches"| node_batch
node_backbones -->|"reads batches"| node_batch
node_objectives -->|"scores batches"| node_batch
node_diffusion -->|"uses model"| node_backbones
node_user -->|"runs training"| node_trainer
node_trainer -->|"fits data"| node_loader
node_trainer -->|"writes checkpoints"| node_checkpoints
node_trainer -->|"returns state"| node_state
node_state -->|"supplies weights"| node_sampling
node_sampling -->|"uses process"| node_diffusion
node_interop -->|"builds components"| node_decoder
node_decoder -->|"supports assembly"| node_interop

click node_sources "https://github.com/ashishkumar4/dew/blob/main/src/dew/data/sources/base.py"
click node_pipeline "https://github.com/ashishkumar4/dew/blob/main/src/dew/data/pipeline.py"
click node_text "https://github.com/ashishkumar4/dew/blob/main/src/dew/data/text.py"
click node_compiler "https://github.com/ashishkumar4/dew/blob/main/src/dew/data/compile.py"
click node_loader "https://github.com/ashishkumar4/dew/blob/main/src/dew/data/dataset.py"
click node_batch "https://github.com/ashishkumar4/dew/blob/main/src/dew/batch.py"
click node_backbones "https://github.com/ashishkumar4/dew/blob/main/src/dew/nn/backbones/causal_transformer.py"
click node_objectives "https://github.com/ashishkumar4/dew/blob/main/src/dew/objectives/base.py"
click node_diffusion "https://github.com/ashishkumar4/dew/blob/main/src/dew/objectives/diffusion/objective.py"
click node_trainer "https://github.com/ashishkumar4/dew/blob/main/src/dew/training/__init__.py"
click node_state "https://github.com/ashishkumar4/dew/blob/main/src/dew/state.py"
click node_checkpoints "https://github.com/ashishkumar4/dew/blob/main/src/dew/checkpoints/__init__.py"
click node_sampling "https://github.com/ashishkumar4/dew/blob/main/src/dew/sampling/__init__.py"
click node_interop "https://github.com/ashishkumar4/dew/blob/main/src/dew/interop/pipeline_assembly.py"
click node_decoder "https://github.com/ashishkumar4/dew/blob/main/src/dew/interop/decoder_parts.py"
click node_serving "https://github.com/ashishkumar4/dew/blob/main/src/dew/inference/serving.py"
click node_evaluation "https://github.com/ashishkumar4/dew/blob/main/src/dew/eval/harness.py"
click node_distributed "https://github.com/ashishkumar4/dew/blob/main/src/dew/distributed/placement.py"
click node_algebra "https://github.com/ashishkumar4/dew/blob/main/src/dew/algebra/program.py"
click node_registry "https://github.com/ashishkumar4/dew/blob/main/src/dew/registry.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_sources,node_pipeline,node_text,node_compiler,node_loader,node_user toneBlue
class node_batch,node_trainer,node_state,node_checkpoints toneAmber
class node_backbones,node_objectives,node_diffusion,node_sampling,node_interop,node_decoder toneMint
class node_serving,node_evaluation,node_distributed,node_algebra,node_registry toneRose
```

## Components

### Data preparation

- [Data sources](https://github.com/ashishkumar4/dew/blob/main/src/dew/data/sources/base.py): `src/dew/data/sources/base.py`
- [Record pipelines](https://github.com/ashishkumar4/dew/blob/main/src/dew/data/pipeline.py): `src/dew/data/pipeline.py`
- [Text processing](https://github.com/ashishkumar4/dew/blob/main/src/dew/data/text.py): `src/dew/data/text.py`
- [Pipeline compiler](https://github.com/ashishkumar4/dew/blob/main/src/dew/data/compile.py): `src/dew/data/compile.py`
- [Data loader](https://github.com/ashishkumar4/dew/blob/main/src/dew/data/dataset.py): `src/dew/data/dataset.py`

### Training core

- [Batch contract](https://github.com/ashishkumar4/dew/blob/main/src/dew/batch.py): `src/dew/batch.py`
- [Trainer](https://github.com/ashishkumar4/dew/blob/main/src/dew/training/__init__.py): `src/dew/training/__init__.py`
- [Training state](https://github.com/ashishkumar4/dew/blob/main/src/dew/state.py): `src/dew/state.py`
- [Checkpoints](https://github.com/ashishkumar4/dew/blob/main/src/dew/checkpoints/__init__.py): `src/dew/checkpoints/__init__.py`

### Models and objectives

- [Model backbones](https://github.com/ashishkumar4/dew/blob/main/src/dew/nn/backbones/causal_transformer.py): `src/dew/nn/backbones/causal_transformer.py`
- [Training objectives](https://github.com/ashishkumar4/dew/blob/main/src/dew/objectives/base.py): `src/dew/objectives/base.py`
- [Diffusion objective](https://github.com/ashishkumar4/dew/blob/main/src/dew/objectives/diffusion/objective.py): `src/dew/objectives/diffusion/objective.py`
- [Sampling](https://github.com/ashishkumar4/dew/blob/main/src/dew/sampling/__init__.py): `src/dew/sampling/__init__.py`
- [Pipeline assembly](https://github.com/ashishkumar4/dew/blob/main/src/dew/interop/pipeline_assembly.py): `src/dew/interop/pipeline_assembly.py`
- [Decoder families](https://github.com/ashishkumar4/dew/blob/main/src/dew/interop/decoder_parts.py): `src/dew/interop/decoder_parts.py`

### Runtime capabilities

- [Inference serving](https://github.com/ashishkumar4/dew/blob/main/src/dew/inference/serving.py): `src/dew/inference/serving.py`
- [Evaluation harness](https://github.com/ashishkumar4/dew/blob/main/src/dew/eval/harness.py): `src/dew/eval/harness.py`
- [Device placement](https://github.com/ashishkumar4/dew/blob/main/src/dew/distributed/placement.py): `src/dew/distributed/placement.py`
- [Kernel algebra](https://github.com/ashishkumar4/dew/blob/main/src/dew/algebra/program.py): `src/dew/algebra/program.py`
- [Model registry](https://github.com/ashishkumar4/dew/blob/main/src/dew/registry.py): `src/dew/registry.py`

### Other components

- **Framework user**

## Connections

- Framework user → Data sources: configures data
- Framework user → Model backbones: chooses model
- Framework user → Diffusion objective: configures objective
- Data sources → Record pipelines: feeds records (evidence: `src/dew/data/pipeline.py`)
- Record pipelines → Text processing: adds text steps (evidence: `src/dew/data/presets.py`)
- Record pipelines → Pipeline compiler: lowers to (evidence: `src/dew/data/pipeline.py`)
- Pipeline compiler → Data loader: builds loader (evidence: `src/dew/data/compile.py`)
- Pipeline compiler → Batch contract: forms batches (evidence: `src/dew/data/compile.py`)
- Model backbones → Batch contract: reads batches (evidence: `src/dew/batch.py`)
- Training objectives → Batch contract: scores batches (evidence: `src/dew/batch.py`)
- Diffusion objective → Model backbones: uses model
- Framework user → Trainer: runs training
- Trainer → Data loader: fits data
- Trainer → Checkpoints: writes checkpoints (evidence: `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 (evidence: `src/dew/interop/pipeline_assembly.py`)
- Decoder families → Pipeline assembly: supports assembly (evidence: `src/dew/interop/decoder_parts.py`)