WKWesley Kambaleinkambale.dev·Jul 29 · 13 min readFeeding the Beast: Data Pipelines with GrainOver the last seven weeks we have built models, written training loops, composed optimizers, bulletproofed our code with Chex, and learned to debug inside JIT-compiled functions. Every one of those ar21N
WKWesley Kambaleinkambale.dev·Jul 18 · 13 min readWhen JIT hides your errors: Debugging JAXLast week, Chex gave us a firewall against shape mismatches and NaN gradients. That firewall catches the bugs that announce themselves; a tensor with the wrong shape, a gradient that explodes into inf20
WKWesley Kambaleinkambale.dev·Jul 12 · 14 min readCatching bugs before they catch youHere's a scenario every ML engineer eventually lives through: you kick off a training run before leaving the office. Three hours later, you check back in. The loss curve looks like a flat line at NaN.00
WKWesley Kambaleinkambale.dev·Jul 1 · 11 min readOptax: Optimizers You Can Compose Like LEGOIn our last article on this blog in March, we built a complete training loop. We used optax.adamw() to create an optimizer, wrapped it in nnx.Optimizer, and watched our model learn. But we barely scra20
AAshiteshinashitesh.me·Jun 11 · 4 min read Scaling Quantum State-Vector Simulation to 36 Qubits on Google Cloud TPU The Problem Simulating quantum circuits on a classical computer is exponentially hard. An n-qubit state vector holds 2ⁿ complex amplitudes — at 30 qubits that's 8 GB, at 33 qubits 64 GB, at 36 qubits 11C
AAshiteshinashitesh.me·Jun 2 · 5 min readSimulating Quantum Computers at Scale: JAX, GPUs, and Cloud TPUsNote: The detailed research paper and full experimental results for this project have already been published. This post provides a high-level summary of the core system design, key benchmarks, and eng10
AAshiteshinashitesh.me·May 31 · 39 min readJAX Quantum Research Suite: A Unified, Hardware-Accelerated, Differentiable Simulator for NISQ-Era Algorithm Research Across GPU and Cloud TPU Clusters We present the JAX Quantum Research Suite, a high-performance, differentiable quantum circuit simulator spanning two co-existing hardware acceleration layers: a GPU division targeting NVIDIA RTX-class30
WKWesley Kambaleinkambale.dev·Mar 4 · 10 min readCustom training loop from scratch in JAXFor the past three weeks, we've been building up to this moment. Week 1 taught us that JAX is fast. Week 2 showed us how to eliminate loops with vmap and compute gradients with grad. Week 3 gave us Fl00
WKWesley Kambaleinkambale.dev·Feb 23 · 11 min readBuilding Neural Networks with Flax NNXOver the past two weeks, we've learned that JAX is fast (jit), that it eliminates loops (vmap), and that it computes gradients automatically (grad). These are powerful primitives. But if you've been f00
WKWesley Kambaleinkambale.dev·Feb 9 · 11 min readTransformations That Change EverythingLast week, we learned that JAX makes code fast through JIT compilation. We took a matrix multiplication from 2 seconds to 0.001 seconds with a single decorator. But speed isn't JAX's only trick. The r00
