Technical blog
Technical articles and experiments
Code-first writeups where the implementation, assumptions, metrics, and limitations are visible.
JAX Inference KV cache Profiling Optimization
What Actually Speeds Up Transformer Inference?
Profiling and optimizing a small autoregressive transformer with JAX, KV caching, batching, graph compilation, and low-bit inference.
Read experiment JAX Flax Optax Transformers Arithmetic
Training a 10M-Parameter Transformer to Learn 3-Digit Arithmetic
A code-first experiment that builds a character-level decoder-only transformer with JAX, Flax, and Optax, then trains it on generated addition, subtraction, multiplication, integer division, and modulo problems.
Read experiment