patrick-kidger

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/

Type annotations and dynamic checking for a tensor's shape, dtype, names, etc.

Ever used asyncio and wished you hadn't? A tiny (~400 lines) event loop for Python.

Code for "Neural Controlled Differential Equations for Irregular Time Series" (Neurips 2020 Spotlight)

Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https://docs.kidger.site/optimistix/

Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax

Differentiable controlled differential equation solvers for PyTorch with GPU support and memory-efficient adjoint backpropagation.

Make posters from Markdown files.

Turn SymPy expressions into trainable JAX expressions.

Differentiable computations of the signature and logsignature transforms, on both CPU and GPU. (ICLR 2021)

Interpolating natural cubic splines. Includes batching, GPU support, support for missing values, evaluating derivatives of the spline, and backpropagation.

Turning SymPy expressions into PyTorch modules.

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

Code for "Deep Signature Transforms" (NeurIPS 2019)

Code for "'Hey, that's not an ODE:' Faster ODE Adjoints via Seminorms" (ICML 2021)

A snappy + easy + pretty TUI debugger for Python.

Code for "Generalised Interpretable Shapelets for Irregular Time Series"