tinyjax

A tiny training engine for transformer LLMs — built as a learning artifact first, working trainer second.

If you know Python and numpy and have wondered what a "framework" like PyTorch or JAX actually does between your model definition and a running GPU kernel, this repo is for you. Every file, function, doc, and notebook is written for that reader.

What tinyjax is

• A minimal, JAX-shaped Python library that trains modern decoder-only transformer language models on CPU and NVIDIA GPU.
• Small enough to hold in one head: about 5–7k lines of library code, an unlimited number of lines of prose.
• Built chapter by chapter. Each chapter introduces one concept (arrays → tracing → jaxpr → autodiff → jit → vmap → transformer → training loop → mixed precision → data parallelism) and ships with code, docs, tests, and a runnable notebook.

What tinyjax is not

• A general-purpose array library (no image models, no encoder-decoder, no sparse ops).
• A production framework (no TPU, no full sharding, no cuDNN fused attention in v1).
• Trying to beat JAX on speed — targeting within 10–20 % of JAX on the same hardware is enough.

Start here

• New to training engines? Read docs/00-what-is-a-training-engine.md.
• Ready to walk through the whole thing? Follow LEARNING_PATH.md top to bottom.
• Ready to contribute? Read CONTRIBUTING.md first.

Install

tinyjax is a normal PEP 621 Python package. Once published on PyPI it will install with any standard tool.

Using uv (fast, recommended):

uv pip install tinyjax

Using pip:

python -m pip install tinyjax

For a local development checkout:

git clone https://github.com/gkatte/tinyjax
cd tinyjax
uv pip install -e ".[dev,docs]"        # or: python -m pip install -e ".[dev,docs]"
pre-commit install
git config --local core.hooksPath .githooks

Requires Python >= 3.12. Runtime dependencies (numpy, ml_dtypes, jaxlib==0.10.2) are pulled automatically. See RELEASING.md for how releases are cut and docs/23-comparing-to-real-jax.md (lands with CH23) for why jaxlib is a pinned dependency.

Status

Under construction. The chapter roadmap lives in LEARNING_PATH.md; the currently-shipped chapters appear there with a check mark.

License

Apache-2.0. See LICENSE and NOTICE.