leafx: a local browser visualizer for Equinox model hierarchies.

Current scope:

• load a serialized model by dragging a file into the browser window
• inspect an operation-style graph left-to-right
• render tensor-carrying edges (including parameter edges such as weight/bias)
• show summary stats (parameter count, bytes, leaf count)
• run the loaded model on a dropped .npy INPUT tensor and inspect OUTPUT shape/values

Quick start

1. Activate your environment:

cd /home/users/armbrust/projects/leafx
source .venv/bin/activate

2. Sync dependencies with uv:

uv sync

3. Start the local server:

uv run uvicorn leafx.app:app --reload --host 127.0.0.1 --port 8000

4. Open in browser:

http://127.0.0.1:8000

Use your own model

Drag one file into the browser window.

Supported file extensions:

• .eqxbundle

The .eqxbundle format bundles factory metadata (to reconstruct the model skeleton) with eqx.tree_serialise_leaves weights in a single zip file.

To generate built-in example files:

cd examples
python example_models.py

This writes .eqxbundle files that can be dragged into the UI.

Example exporter:

import jax.random as jr

from example_models import make_mlp
from leafx.model_loading import save_model_bundle

model = make_mlp(in_size=32, out_size=4, width=64, depth=3, seed=0)
save_model_bundle(
	model,
	"model.eqxbundle",
	factory_spec="example_models:make_mlp",
	factory_kwargs={"in_size": 32, "out_size": 4, "width": 64, "depth": 3, "seed": 0},
)

Run inference from browser input

After loading a model, drop a .npy tensor directly onto the INPUT node. leafx will run a forward pass on the loaded model and populate the OUTPUT node with inferred output shape and a heatmap preview of output values.

Notes for large models

• operation nodes are laid out left-to-right by depth
• tensor labels are drawn on edges; parameter edges are highlighted
• use browser zoom/pan for navigation
• summary shows global parameter and memory size totals

Next extensions

• parameter-mask extraction and save/load
• per-subtree statistics (norms, sparsity, histograms)
• architecture editing and round-trip export