Metadata-Version: 2.4
Name: tinyjax
Version: 0.1.0.dev0
Summary: A minimal JAX-like training engine for transformer LLMs — built as a learning artifact.
Author: Nikhil Gajendrakumar
License: Apache-2.0
Project-URL: Homepage, https://github.com/gkatte/tinyjax
Project-URL: Documentation, https://github.com/gkatte/tinyjax/blob/main/docs/00-what-is-a-training-engine.md
Project-URL: Repository, https://github.com/gkatte/tinyjax.git
Project-URL: Issues, https://github.com/gkatte/tinyjax/issues
Project-URL: Changelog, https://github.com/gkatte/tinyjax/blob/main/CHANGELOG.md
Keywords: jax,training,transformer,llm,autodiff,learning
Classifier: Development Status :: 2 - Pre-Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Education
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: numpy>=2.1
Requires-Dist: ml_dtypes>=0.5.0
Requires-Dist: jaxlib==0.10.2
Provides-Extra: dev
Requires-Dist: pytest>=8; extra == "dev"
Requires-Dist: pytest-xdist>=3; extra == "dev"
Requires-Dist: pre-commit>=4; extra == "dev"
Requires-Dist: ruff>=0.15; extra == "dev"
Requires-Dist: build>=1.2; extra == "dev"
Requires-Dist: twine>=5; extra == "dev"
Requires-Dist: jupyter>=1; extra == "dev"
Requires-Dist: nbconvert>=7; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx>=8; extra == "docs"
Requires-Dist: myst-parser>=4; extra == "docs"
Requires-Dist: alabaster>=1; extra == "docs"
Requires-Dist: sphinxcontrib-mermaid>=1; extra == "docs"
Provides-Extra: gpu
Requires-Dist: jaxlib==0.10.2; extra == "gpu"
Dynamic: license-file

# 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`](docs/00-what-is-a-training-engine.md).
- Ready to walk through the whole thing? Follow [`LEARNING_PATH.md`](LEARNING_PATH.md) top to bottom.
- Ready to contribute? Read [`CONTRIBUTING.md`](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):

```bash
uv pip install tinyjax
```

Using **pip**:

```bash
python -m pip install tinyjax
```

For a local development checkout:

```bash
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`](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`](LEARNING_PATH.md); the currently-shipped chapters appear there with a check mark.

## License

Apache-2.0. See [`LICENSE`](LICENSE) and [`NOTICE`](NOTICE).
