Sitemap

A list of all the posts and pages found on the site. For you robots out there is an XML version available for digesting as well.

Pages

Posts

Trainable kernels in Jax

It is not difficult to write up a kernel \(k_{\theta}(x,y)\) in Jax and there are many ways to do so. Here I’ll offer my method, which is to just maintain a class module with self-contained trainable parameters, just as one might write a custom neural network layer. Therefore, we don’t have to carry around parameters elsewhere, and if we want to keep them fixed, we can.

portfolio

publications

talks

teaching

tutorials

Continuous Normalizing Flow

Continuous-time density transformation with Diffrax, Equinox vector fields, and the instantaneous change-of-variables formula.

Flow Matching

Simulation-free vector-field training along conditional probability paths, implemented with JAX, Equinox, Optax, and Diffrax.

Normalizing Flow

RealNVP coupling layers, change of variables, maximum likelihood training, and sampling in JAX and Equinox.