diffrax

收录理由

连续时间模型在机器学习里越来越常见，Diffrax 是 JAX 生态中让这类模型真正落地的关键工具。它把常微分、随机微分和受控微分方程的求解器都整合在一起，从 Tsit5、Dopri8 到辛积分器和隐式方法都有，选择相当丰富。代码风格贴合 JAX 习惯：可以对批量和积分区间做 vmapp，状态用 PyTree 表示，还能拿到步与步之间的稠密解，放进训练循环里很顺手。反向传播支持多种伴随方法，梯度能干净地穿过求解过程。做神经微分方程或物理信息模型的团队，不用自己维护求解器，用它就能获得可微分的数值积分，省心也可靠。