-
diffrax
Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/
-
InfluxDB
Power Real-Time Data Analytics at Scale. Get real-time insights from all types of time series data with InfluxDB. Ingest, query, and analyze billions of data points in real-time with unbounded cardinality.
Haha! So until recently I would said "unfortunately they work really badly on GPUs and this sucks". As it turns out, however, a large part of this was actually just the overhead of the Python interpreter inside libraries like torchdiffeq (which has been one of the go-to libraries for working with neural ODEs over the past few years). Whilst it depends a lot on the exact problem, I have sometimes observed dramatically better performance (and GPU utilisation etc.) with the new Diffrax library I mention in the main post. As this is jit-compiled using JAX, the overhead of the Python interpreter is no longer present.