Blog posts

Today we’re going to implement the aforementioned paper in code. Reason 1: They did not release code. Reason 2: I am using their method in my research so I had to code it up. Reason 2: Save other people’s time? Reason 3: Noooo codddieeeee.

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.