programming

JAX

Improved batch norm with hijax?

Originally Publish:  February 11, 2026 Note: this was published before hijax became somewhat established in JAX, I will investigate the new changes and put out a new blog on them. Stateful operations in JAX, such as batchnorm, can sometimes be annoying to work with. As a DSL that relies on

JAX

Numpy Random Interface in JAX

Originally published: December 13, 2025 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