jepax v0: an implementation of IJEPA training in JAX/Equinox
Published:
Blog posts
2026
Improved batch norm with hijax?
Published:
Stateful operations in JAX, such as batchnorm, can sometimes be annoying to work with. As a DSL that relies on pure functions, it requires the usual functional approach of passing around a different state object to all the functions, changing $f :: x \to y$ to $f :: (x, state) \to (y, state)$. While this may not sound like much work, it can be quite tedious to manage with large and complex machine learning network workflows. Part of this complexity comes from the fact that we often want to use vmap to batch, rather than having a batch dimension for everything. Although different libraries in JAX implement stateful operations in different ways, they have a common mechanism for batchnorm, all of them keep track of a vmap named axis, then do collective operation like pmean to get the stats across the vmapped batch. Here, we won’t be looking to simplify that dimension of batchnorm, but the state management.
2025
Numpy Random Interface in JAX
Published:
In this blog, we’ll be looking at one of the “sharp bits” of JAX: pseudorandomness, and see if we can’t smooth it out a little. JAX is an accelerated linear algebra library that forms the backbone of some modern python based scientific and machine learning libraries (things like Google’s Gemini were built on JAX). Personally I’m a fan.
Can Quantization Reduce Separability?
Published:
I was recently experimenting with the linear separability of embeddings, and showed a graph which indicated the embeddings were substantially more linearly separable than the quantized embeddings. One person commented that this didn’t make sense, and the quantized vectors should be as separable as the non-quantized vectors. That didn’t seem right to me, but on the spot I couldn’t think of a trivial counter example. In this blog, I present two.
Why The AI Bubble Will Never Pop
Published:
As the AI hype train has progressed and consumed upwards of 35% of the US economy (which surely is not a problem), a growing sentiment has emerged: that this is a bubble. This is probably now even the dominant opinion: that given the scale of investment, the large valuations, the lack of returns, and the slow progress, this constitutes a bubble. And bubbles always pop.