Software

I’m an author, developer, or major contributor to several pieces of software. Many of these are to do with scientific computing and scientific machine learning; in particular differential equation solvers. I am proudest of my work in the JAX ecosystem.

See also my GitHub page.

JAX

Equinox

For building neural networks (or in general any parameterised function). PyTorch-like API, fully compatible with native JAX, with no new concepts you have to learn.

Diffrax

An advanced suite of numerical differential equation solvers. Very efficient; features ODE/SDE/CDE solvers, high-order solvers, implicit solvers, dense solutions, adjoint methods, etc.

Lineax

Linear solvers and linear least squares, with fast compile times and efficient autodiff.

Optimistix

Root-finding, minimisation, nonlinear least squares, fixed points.

jaxtyping

Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees.

sympy2jax

Build your physics-informed model in SymPy, perform arbitrary symbolic manipulations on it, then convert it to JAX and train it via gradient descent.

PyTorch

TorchTyping

Rich type annotations for a tensor’s shape, dtype, etc. Includes optional runtime type checking.

torchdiffeq

Ordinary differential equation (ODE) solvers.

torchsde

Stochastic differential equation (SDE) solvers.

torchcde

Controlled differential eqution (CDE) solvers. In particular useful for building Neural Controlled Differential Equations on time series.

sympytorch

Turn SymPy expressions into PyTorch modules and back again. Train your SymPy expressions by gradient descent.

torchcubicspline

Implements natural cubic splines.

Signatory

Differentiable computations of the signature and logsignature transforms. ICLR 2021 paper.

Julia

FromFile.jl

An improved import/include system for Julia. Makes your files self-contained and easier to understand.

Other

MkPosters

Write academic posters in Markdown, style them with CSS, save them to PDF. No wrestling with LaTeX.

typst_pyimage

Typst extension, adding support for generating figures using inline Python code.