← back to the tree

ml-eng-lab

Twenty-one ML tasks as executed notebooks — GNNs to diffusion to DPO — built on NNx.

21 task folders29 executed notebooksTier-A CI

make atlas-setup && make atlas-up

ML Eng Lab — notebooks, systems, reproducibility — set in a glowing field of networks and gradients.
Notebooks, systems, reproducibility. Brand banner, 2026 — ml-eng-lab repository.

Reading about a machine-learning technique and having actually run it are different things. ml-eng-lab is where each one gets run. Twenty-one active task folders hold twenty-nine executed notebooks, every folder named [task]-[dataset]-[model]-[framework] so the catalog reads at a glance: diffusion-mnist-ddpm-pytorch, preference_alignment-toy-dpo-pytorch, node_classification-reddit-gnn-pyg. Each folder is self-contained — the notebook, a README that explains the idea next to the code, and its own data and run directories — and each notebook is saved with its outputs, so the results on GitHub are the results that actually happened.

The folders share a PyTorch toolkit, NNx (published as thekaveh-nnx on PyPI), which owns the training-loop, dataset, and plotting primitives so each notebook can stay about the model. Library and lab co-evolve: a new task lands its nnx additions upstream first, then pins the released version back here — so every experiment exercises the library’s public surface exactly as an outside user would.

Highlights

  • From scratch to framework — the same MNIST feed-forward classifier built twice, once in pure NumPy and once in PyTorch via nnx, so what the framework buys you is visible, plus tabular classification on Iris and regression on Diabetes against scikit-learn baselines
  • A nine-notebook GNN study — node classification on the Reddit2 graph in three phases (dataset exploration, model selection, main training and evaluation) comparing GraphConv, GraphSAGE, and GAT in PyTorch Geometric, with link prediction and community detection on the Karate Club graph alongside
  • Generative and language modeling — a decoder-only transformer trained on TinyShakespeare with a BPE tokenizer and a full sampling stack, a DDPM diffusion denoiser on MNIST built without a U-Net, and DPO preference alignment training a policy against a frozen reference model
  • Self-supervision and sparsity — I-JEPA pretraining of a ViT on Fashion-MNIST with an EMA target network and a linear probe, and a Mixture-of-Experts classifier with four experts and top-2 routing
  • A model-efficiency suite — LoRA versus DoRA adapters for cross-task adaptation from MNIST to Fashion-MNIST, post-training quantization and quantization-aware training with torchao, a magnitude-pruning sparsity sweep, born-again knowledge distillation, and Net2Net model surgery
  • Classical baselines everywhere — PCA against an autoencoder, KMeans on raw features against KMeans on a learned latent, VADER against a neural sentiment model, Louvain against a GNN, and a spaCy bag-of-words MLP for four-topic text classification, so every neural result has something to beat
  • Tiered re-execution, verified in CI — cheap notebooks re-run under papermill on every pull request; heavier sweeps and the historical GPU training runs execute smoke-truncated so their committed outputs stay authentic
  • A preserved archive — twenty-two CodeXGLUE code-summarization experiments with HuggingFace transformers from 2023, kept read-only

Stack & Architecture

Python and PyTorch throughout, with PyTorch Geometric for the graph work, torchao for quantization, and spaCy, NLTK, and scikit-learn where the classical contrasts call for them. The notebooks are the deliverable; nnx is installed from PyPI rather than vendored, and a pytest suite guards its consumed surface alongside a notebook verifier. Four runtimes are documented and maintained: a vendored JupyterHub stack as the primary, local Docker, a plain venv, and a devcontainer that gives Codespaces the same pinned environment in one click.

Status

An active portfolio, MIT licensed, with a roadmap of queued tasks — CIFAR-10 ResNet, XGBoost on Titanic, a DistilBERT fine-tune — and a generated documentation site built from the in-repo docs.