Skip to content
HaozheZhang6Public

About

Publication-quality, hierarchical figures of model architectures & training methods — an agent reads the code, then draws the figures

Topics

Resources

Contributing

Stars

3 stars

Watchers

0 watching

Forks

Latest commit

 

History

37 Commits

Folders and files

Repository files navigation

archscope

License: MIT Python 3.10+

Publication-quality, hierarchical figures of model architectures and training methods — built for the workflow: an agent reads the code, then draws the figures.

No auto-layout, no op-soup. A figure is a small Python script; the library supplies NeurIPS-style components, measurement-based layout, perpendicular edge routing, and vector math — you (or your coding agent) supply the understanding.

Pre-LN transformer block

# one call: stack in dataflow order, auto-id from labels, auto-connect with arrows
d.flow([
    IOLabel("x  (B, T, 768)"),
    Block("LayerNorm", kind="norm"),
    Block("MLP", kind="ffn", sub="768 → 3072 · GELU · 3072 → 768", src="model.py:42"),
    OpDot("+", id="add"),
    IOLabel("y  (B, T, 768)"),
], dir="up", x=240, y=120)

d.auto_legend(160, 86)        # legend built from what was used
assert not d.check()          # structured overlap report — no need to open the PNG

No per-box coordinates, no manual chain, no hand-written legend. Place a second column relative to the first with d.place(el, right_of="enc", gap=80).

What it draws

  1. Block anatomy — what a tensor passes through (attention / FFN / RMSNorm / VAE / …), decomposed level by level (system → model → block → sublayer), every claim annotated with a file:line source pointer.
  2. Tensor flow — shape-annotated dataflow at every level of the hierarchy.
  3. Training methods — the genre almost no tool covers: objectives as storyboards (token rows, attention-mask grids, noise schedules, loss formulas as outlined vector math).

Concepts the reader already knows ("I know what attention is") collapse into a ✓-chip instead of being expanded — decomposition depth is a per-reader choice.

Install & run

pip install -e .            # only dependency: matplotlib
brew install resvg          # optional: PNG export (SVG always works)

python examples/quickstart.py          # ~40-line API tour

Outputs land in out/ as self-contained SVG (fonts as text, math as outlined paths — survives every converter) plus 4× PNG. Full index: examples/.

Gallery

Qwen3 decoder layer Qwen3 decoder layer, as implemented — fused GQA shapes, per-head QK-RMSNorm, two-branch SwiGLU; every box cites nano-vllm file:line. GPT vs BERT training storyboards Training objectives as storyboards — GPT next-token vs BERT masked-LM; the attention-mask grids are computed from the rule (j ≤ i), never hand-drawn.
Flow matching storyboard Flow matching (rectified flow) — the SD3/FLUX/Wan objective: one straight path, regress its velocity, integrate it back. LingBot-VA DiT block Case study: LingBot-VA — 8 figures from an agent reading real code; per-token AdaLN block shown here. Full gallery: out/lingbot_va.

Case study: LingBot-VA

examples/lingbot_va/ + out/lingbot_va/ — an 8-figure set for Ant Group's LingBot-VA video-action world model, drawn by an agent that read the code and the paper. It surfaced a real finding: the paper describes a dual-stream MoT (video d=3072 / action d=768), while the released code ships a single shared-weight stream routed purely by the attention mask — including the fossil evidence in the config lists.

MoT: paper vs released code

Design rules the library enforces

  • Edges meet boxes perpendicular to the side they touch, with entry/exit stubs (the default router guarantees it; raw-point endpoints take a_side/b_side).
  • Repeated blocks are card stacks (RepeatStack): solid front card, dashed ghost outlines behind; edges attach to the stack envelope and never cross the ghosts.
  • Computed truth: mask grids / schedules are generated by re-implementing the model's actual rules inside the figure script, never hand-transcribed (see examples/training_methods.py and the LingBot fig 5).
  • Semantic palette: one pastel family per component kind (attention, FFN, norm, conditioning, VAE, …), auto-legend; modality stripes for multi-modal models.
  • Glyph-safe text (no .LastResort tofu), DejaVu-metric measurement so boxes never underflow, KaTeX-free math via matplotlib mathtext → SVG paths.

Authoring rules (for agents and humans)

  • Verify before drawing: every node/edge fact comes from code you read (annotate file:line) or the paper (annotate §/page). Paper-vs-code conflicts are content, not noise — draw both and say so.
  • Read the rendered output before calling a figure done. Overlaps are bugs.
  • A complex model is several figures linked by → Fig N pointers, not one mural.

Status & roadmap

Early but used in anger (the LingBot-VA set is a real work product, and the library is exercised daily by a coding agent). Planned: an MoE / multi-modal example set, an optional torch tracer for shape cross-checks, a JSON spec layer once the grammar stabilizes, PDF export. Issues and PRs welcome.

Citation

@software{archscope,
  author = {Zhang, Haozhe},
  title  = {archscope: publication-quality figures of model architectures and training methods},
  url    = {https://github.com/HaozheZhang6/archscope},
  year   = {2026}
}

License

MIT. Built with Claude Code driving the library — which is also its intended user.

About

Publication-quality, hierarchical figures of model architectures & training methods — an agent reads the code, then draws the figures

Topics

Resources

Contributing

Stars

3 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages