Repolode
提交项目
patrick-kidger

patrick-kidger/diffrax

基于 JAX 的数值微分方程求解库，支持 ODE、SDE 和 CDE，可自动微分并运行于 GPU。

在 GitHub 查看
Star
2,122
Fork
202
未关闭 Issue
265
最近提交
2个月前

这是什么？

Diffrax 是基于 JAX 构建的数值微分方程求解库，覆盖常微分方程（ODE）、随机微分方程（SDE）和受控微分方程（CDE）。三类方程采用统一的内部结构求解，而不是分别实现。

库中内置多种求解器，包括用于显式积分的 Tsit5、Dopri5、Dopri8，以及辛方法和隐式方法。积分过程可以对所有输入（包括积分区间）做 vmap，状态可以是任意 PyTree。它还支持稠密输出，并提供多种伴随方法用于反向传播。

适合需要可微分仿真的场景，例如神经微分方程。通过 pip 安装，要求 Python 3.10 及以上版本。

核心功能

  • 01用统一的内部结构求解 ODE、SDE 和 CDE。
  • 02内置 Tsit5、Dopri8 等显式、辛方法和隐式求解器。
  • 03求解过程支持 vmap，包括积分区间。
  • 04状态可以是 PyTree，并支持稠密输出。
  • 05提供多种伴随方法用于反向传播。

适用场景

  • 训练需要通过求解器计算梯度的神经微分方程模型。
  • 在 GPU 上批量仿真动力系统。
  • 使用 SDE 求解器模拟随机过程。

快速上手

  1. 1

    使用 pip 安装该库

    pip install diffrax
  2. 2

    定义 ODE 项和求解器，然后调用 diffeqsolve（参见 README 中的快速示例）

使用前须知

  • 需要 Python 3.10 或更高版本。
  • 基于 JAX 构建，因此使用前需要正确配置 JAX 环境。

收录于专题

GPU 加速的 AI 训练与推理工具

标签

#differential-equations#jax#scientific-computing#neural-differential-equations#ode-solver#gpu#deep-learning#dynamical-systems#equinox#machine-learning

常见问题

diffrax 是什么？
基于 JAX 的数值微分方程求解库，支持 ODE、SDE 和 CDE，可自动微分并运行于 GPU。 Diffrax 是基于 JAX 构建的数值微分方程求解库，覆盖常微分方程（ODE）、随机微分方程（SDE）和受控微分方程（CDE）。三类方程采用统一的内部结构求解，而不是分别实现。
diffrax 可以免费使用吗？用的什么许可证？
diffrax 以 Apache-2.0 许可证开源。用于商业产品前，请先确认许可证条款。
diffrax 还在维护吗？
半年内有更新。diffrax 最近一次提交是在 2026年9月6日。
diffrax 适合用来做什么？
训练需要通过求解器计算梯度的神经微分方程模型。 在 GPU 上批量仿真动力系统。 使用 SDE 求解器模拟随机过程。
diffrax 怎么安装和运行？
使用 pip 安装该库: pip install diffrax 定义 ODE 项和求解器，然后调用 diffeqsolve（参见 README 中的快速示例）
在哪里获取 diffrax？
源代码和安装说明在 GitHub 上：https://github.com/patrick-kidger/diffrax

概要由 AI 根据项目 README 生成；GitHub 数据每日更新。

相关项目