GitHub repository: lockwo/diffrax_extensions
Primary language: Python
Description: Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable.