@kambale

Custom training loop from scratch in JAX

Mar 4 · 10 min read · For 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 Fl

Building Neural Networks with Flax NNX

Feb 23 · 11 min read · Over 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 f

Transformations That Change Everything

Feb 9 · 10 min read · Last 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 real power of JAX lies in its transformations; func...

Why JAX? The NumPy You Know, But Faster

Feb 2 · 9 min read · If you've been doing machine learning in Python for any length of time, you've written code like this: import numpy as np x = np.random.randn(1000, 1000) y = np.random.randn(1000, 1000) result = np.dot(x, y) NumPy is comfortable. It's the first thi...

The reverse turing test: We must now prove we are “dumb” to beat AI

Dec 10, 2025 · 4 min read · In the dystopian logic of the digital age, a new anxiety has gripped the writing world. From university lecture halls in Makerere to newsrooms in Kampala, humans are facing a pressure that would have seemed laughable just two years ago: the pressure ...