MNIST Runtime Lab

mnist-runtime-lab is a notebook-first project for comparing how the same small vision workload behaves across different execution models.

The main notebook walks through:

• Pure Python loops as an interpreter-bound baseline
• NumPy vectorization on CPU
• PyTorch on CPU and CUDA, including optional profiling and dynamic quantization
• JAX on CPU/CUDA with jit, separating compile time from steady-state execution
• A larger matrix multiplication microbenchmark to make hardware effects easier to see

This project is designed to make performance tradeoffs legible:

• why scalar Python is slow for numerical work
• why vectorized kernels help so much on CPU
• what changes when tensors move to GPU
• how eager execution differs from JAX compilation
• when quantization helps and when it does not

.
|-- mnist_perf_comparison_notebook.ipynb
|-- requirements.txt
|-- data/                # local dataset cache, ignored by git
|-- torch_traces/        # local profiler output, ignored by git
`-- jax_traces/          # local profiler output, ignored by git

python -m venv .venv
.venv\Scripts\Activate.ps1
python -m pip install --upgrade pip
pip install -r requirements.txt

Then open:

jupyter notebook mnist_perf_comparison_notebook.ipynb

• torch and torchvision are included in requirements.txt for the PyTorch sections.
• JAX support is optional in the notebook and installation is platform-specific. If you want to run the JAX section, install the appropriate JAX build for your machine after the base requirements.
• The notebook downloads FashionMNIST into ./data when needed.