A ~20M-parameter decoder-only transformer trained on TinyStories, built from scratch with JAX and Flax NNX. Source: llm-with-Jax.