Name: JAX Talk: Diffrax
Start: 2022-05-20T09:00:00-07:00
End: 2022-05-20T10:30:00-07:00

In this session, we'll discuss “Diffrax”, a JAX-based suite of ordinary/stochastic/controlled differential equation solvers, roughly analogous to the existing torchdiffeq/torchsde/DifferentialEquations.jl packages, which are available in the PyTorch and JAX ecosystems.

Highlights include:

* High performance, e.g. 200x speedup over torchdiffeq and with similar performance to DifferentialEquations.jl.
* Numerous features: high-order solvers, implicit solvers, dense solutions, multiple adjoints methods, etc.
* Integrates directly with JAX: jit/grad/vmap/etc. all work as normal.
* Easily extensible with custom ops (solvers etc.); includes the ability to handle the stepping yourself if writing a differentiable simulator yourself.
* For the numerical analyst: ODEs/SDEs/etc. are all solved in a unified way.

Sanyam Bhutani

Cristian Garcia

Events

JAX Global Meetup

Technology

Artificial Intelligence

Data Science

Computer Programming

Predictive Analytics

Data Visualization

Data Analytics

New Technology

JavaScript

Software Development

Python

Machine Learning

High Performance Computing

Deep Learning

JAX Talk: Diffrax

Online event

Share

JAX Global Meetup

JAX Talk: Diffrax

JAX Global Meetup

Details

Members are also interested in