Whitebox Studio

A visual builder for deep-learning architectures — with native observability and a bidirectional Python DSL.

Wire neural blocks on a canvas, watch a real PyTorch model train, and see inside it — activations, weights, gradients and attention streaming into live plots at every point of the graph.

Quick start · Features · How it works · Roadmap · Contributing

---

What is this?

Whitebox Studio is a single-window engineering tool for defining, training, testing and analysing deep-learning architectures — from an MLP up to a nanoGPT-style transformer. It is built on three principles:

• 🔍 Observability-first. Every wire is inspectable; every node is probeable with one click; the plots live inside the editor, not in a separate notebook.
• 🔁 One IR, two views. A single canonical intermediate representation. The visual canvas and the Python DSL are two editable, synchronised projections of the same graph.
• 🪟 Transparent where it can be, honest where it can't. Formula nodes are white-box equations; neural blocks open down to their math with opaque weights. For big models the value is in scale/data/compute — so the tool defines and exports them, and trains small versions you can actually watch.

Everything you see streaming is real torch — no mocked data. The char-level LM example trains and decodes Shakespeare on a laptop CPU in seconds.

Whitebox Studio is the running implementation of the design spec in specifica-tecnica-whitebox-dl.md, wearing the Neural Shape design system.

Screenshots

Drop your captures into docs/ — see Contributing. Suggested: docs/editor.png, docs/training.png, docs/console.png.

Edit	Train	Inspect
Wire the graph; live shape inference on every port.	Loss curve + activation / ‖w‖ / ‖∇‖ / attention probes streaming.	Step-by-step decoding with the next-token distribution.

Quick start

You need Python ≥ 3.10 and Node ≥ 18.

git clone https://github.com/TheSmallPixel/NeuralShape.git
cd NeuralShape
./run.sh

run.sh creates the engine virtualenv (installs torch, fastapi, …), installs the studio dependencies, then starts:

• the engine on http://127.0.0.1:8000
• the studio on http://localhost:5173 ← open this

Run the two halves manually

# engine (Python sidecar)
cd engine
python3 -m venv .venv && .venv/bin/pip install -r requirements.txt
.venv/bin/python -m whitebox_engine.server

# studio (frontend) — in a second terminal
cd apps/studio
npm install && npm run dev

The Vite dev server proxies /api and /ws to the engine, so the UI uses same-origin URLs.

First five minutes

1. The studio opens on the mini-GPT example (char-level LM on bundled Shakespeare). Hit Train and watch the loss fall and the probe grid fill.
2. Right-click a node (or the coral pin) to probe it — its activation / weight / gradient / attention plots appear in the dock.
3. Switch to Test → type a prompt → Generate to decode token-by-token with the next-token distribution shown at each step.
4. Export → PyTorch emits a standalone, runnable nn.Module; Import DSL parses edited DSL straight back onto the canvas.
5. Open the Inspector (nothing selected) to switch datasets, edit the training config, and Save / Open the project.

Features

Build

• Drag nodes from a typed palette; drag from a port to wire, with live type/shape validation (compatible ports glow, incompatible ones go red, cycles rejected).
• Click a wire to select it; Delete or right-click to remove. Esc cancels a wire mid-drag.
• Symbolic shape inference across the whole graph, shown inline on every port ([B, t, 64]).
• Resizable, dockable zones; movable/resizable floating panels — never a modal over your work.

Train (real PyTorch, CPU by default)

• Optimizer + LR schedule (warmup → cosine) + gradient clipping; live start / pause / step / stop and hot-swappable LR.
• Two trainable starter architectures: mini-GPT (transformer char-LM) and an MLP classifier.

Observe (the heart of the product)

• One-click probes on any port → an in-node mini-scope plus a full tile in the dock.
• Plot types: loss/metric curves, activation & weight ‖w‖ histograms, gradient ‖∇‖ distributions, attention heatmaps, sample-axis series, raw tensor inspector.
• Two time axes: a sample cursor and an iteration scrubber (live or scrubbed-to-step).
• Diagnostics: NaN/Inf detection, exploding/vanishing-gradient warnings — surfaced in a Console with an error badge.

Test & inference

• Single forward on a chosen input with per-node tensor inspection (shape, μ/σ, min/max, preview).
• Autoregressive decoding for language models with temperature / top-k / top-p and the next-token distribution shown per step.

Code & I/O

• Round-trip DSL: graph → DSL → graph is the identity (parsed with Python's ast, never eval).
• Export to a standalone runnable PyTorch nn.Module and to the Whitebox DSL.
• Save / open projects as portable IR JSON.

The node library

Group	Nodes
Data / IO	Input, Target
Neural	Embedding, Positional encoding (learned/sinusoidal/RoPE), Linear, Activation, Dropout, LayerNorm, RMSNorm, Add (residual), FFN/MLP (incl. SwiGLU)
Attention	Multi-head attention (causal), Transformer block (Stack ×n)
Recurrent	LSTM, GRU
Tricks	Stop-gradient
Loss	Cross-entropy, MSE
Optimizer	AdamW, Adam, SGD

Color encodes meaning, never decoration: teal = data/neural, coral = time/recurrence/gradient-tricks, violet = metrics, gold = loss.

How it works

┌─────────── Renderer (browser) ───────────┐        ┌──── Engine (Python sidecar) ────┐
│  React + Zustand                          │  WS    │  FastAPI + WebSocket            │
│  • IR (source of truth) + shape inference │ <────> │  • IR → torch.nn compiler       │
│  • 5-zone editor (design system)          │  JSON  │  • training runtime + probes    │
│  • live probe plots / dock                │        │  • forward / decode / codegen   │
└───────────────────────────────────────────┘        └─────────────────────────────────┘

• The IR lives in the renderer — it is the single source of truth. The engine is keyed by graph signature and holds only derived artifacts (the compiled function, weights, optimizer state, probe buffers).
• The compiler builds the torch model by tracing an example batch, so every layer is sized to its real input shape; the same per-node forward path is reused for training, inference and codegen.
• Probes are throttled and sub-sampled at plot resolution with latest-wins delivery — the full tensors stay in the engine.

Project structure

apps/studio/                  # Vite + React + TS frontend (the renderer)
  src/ir/                      # IR types + the engine WebSocket protocol
  src/nodes/                   # node registry (vocabulary), shape inference, validation
  src/store/                   # Zustand store + example graphs
  src/components/ds/           # design-system primitives (copied from the DS package)
  src/zones/                   # Toolbar · Palette · Canvas · Inspector · Dock · panels
  src/plots/                   # SVG plot renderers
engine/whitebox_engine/        # Python engine
  ir.py compiler.py training.py inference.py codegen.py dslparse.py datasets.py probes.py server.py
engine/smoke.py                # standalone engine test (no server needed)
Neural Shape Design System/    # the brand: tokens, components, UI kit (source of truth)
run.sh                         # one-command dev launcher

Built-in datasets

id	task	used by
tiny-shakespeare	char-level language modelling (bundled text)	mini-GPT
two-moons	2-D binary classification	MLP
spiral	3-class spiral classification	—
sine	1-D regression	—

Roadmap

Working today

• Visual graph editor with five resizable zones
• Canonical IR + symbolic shape inference + connection validation
• IR → torch.nn compilation and training (mini-GPT, MLP)
• Live observability: probes, loss curves, histograms, attention, two time axes, diagnostics console
• Test / manual mode: per-node tensor inspection, autoregressive decoding
• Bidirectional DSL round-trip + runnable PyTorch export
• Project save / open (IR JSON)

Planned

• Undo / redo
• A/B run comparison (overlay two runs)
• Custom-block authoring (select sub-graph → group → save to palette)
• JEPA / SSL primitives (EMA target coupling, masking pipelines)
• Weight checkpoints to disk; ONNX export
• GPU / distributed training

Requirements

• Python ≥ 3.10 (verified on 3.14 with torch 2.11). The engine is CPU-by-default; GPU is transparent where available.
• Node ≥ 18.

Tech stack

Frontend — React 18, TypeScript, Vite, Zustand, SVG plots, the Neural Shape design system (IBM Plex, dark-first). Engine — Python, PyTorch, FastAPI + WebSockets, NumPy, SymPy.

Contributing

Contributions are very welcome — this is an alpha and there's a lot of fun work on the roadmap.

1. Fork and clone, then ./run.sh to get both halves running.
2. Frontend: cd apps/studio && npm run dev; type-check with npx tsc -b --noEmit.
3. Engine: edit engine/whitebox_engine/, sanity-check with engine/.venv/bin/python engine/smoke.py.
4. Keep the node vocabulary in sync: apps/studio/src/nodes/registry.ts (frontend) mirrors the builders in engine/whitebox_engine/compiler.py.
5. Follow the design system — color is meaning, numbers are mono/tabular, no emoji in product chrome, respect prefers-reduced-motion.

Open an issue to discuss larger changes first. Bug reports with a saved project JSON attached are gold.

To add screenshots referenced above, capture the editor at ~1440×900 and commit them under docs/.

Acknowledgements

• The Neural Shape design system and the Whitebox Studio product spec that this implements.
• Built on PyTorch, FastAPI, React and Vite.
• The mini-GPT example is inspired by Andrej Karpathy's nanoGPT / char-RNN lineage.

License

MIT © TheSmallPixel