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

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.

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.

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.

jaxtyping

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

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.