patrick-kidger/diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

What it solves

Diffrax는 JAX 생태계 내에서 미분 방정식—특히 일반 미분 방정식 (ODE), 확률 미분 방정식 (SDE), 제어 미분 방정식 (CDE)—을 수치적으로 풀 수 있는 방법을 제공합니다. 이 솔버들은 자동 미분이 가능하고 GPU에서 실행될 수 있어 물리 기반 모델을 머신러닝과 통합하는 데 필수적입니다.

How it works

JAX 위에 구축된 이 라이브러리는 통합된 내부 구조를 사용해 다양한 유형의 방정식을 처리합니다. Tsit5와 Dopri8 같은 다양한 솔버를 지원하고, 상태를 PyTree 형태로 표현할 수 있으며, 역전파를 위한 여러 adjoint 방법을 제공합니다. JAX 기반이기 때문에 적분 구간을 포함한 프로세스의 거의 모든 부분이 vmappable합니다.

Who it’s for

과학 컴퓨팅 및 AI 분야에서 연구 및 개발을 하는 연구자와 개발자를 위해 설계되었으며, 특히 신경 미분 방정식을 구현하거나 고성능·미분 가능한 수치 솔버가 필요한 경우에 적합합니다.

Highlights

ODE, SDE, CDE 솔버 지원.
완전 자동 미분 및 GPU 호환.
라이브러리 전반에 걸친 vmappable 연산.
신경 미분 방정식 지원.
심플렉틱 및 암시적 옵션을 포함한 다양한 솔버 제공.