Paquete instalable con pip que fija un conjunto probado de bibliotecas de JAX (Flax, Optax, Orbax, Chex y Grain) para IA y computación numérica.
patrick-kidger/diffrax
Solvers numéricos de ecuaciones diferenciales para JAX, con soporte para ODE, SDE y CDE, autodiferenciación y ejecución en GPU.
- Estrellas
- 2,122
- Forks
- 202
- Issues abiertos
- 265
- Último commit
- hace 2 meses
¿Qué es?
Diffrax es una biblioteca de solvers numéricos de ecuaciones diferenciales basada en JAX. Resuelve ecuaciones diferenciales ordinarias (ODE), ecuaciones diferenciales estocásticas (SDE) y ecuaciones diferenciales controladas (CDE) mediante una estructura interna unificada, de modo que las tres se resuelven de la misma manera.
La biblioteca incluye numerosos solvers, como Tsit5, Dopri5 y Dopri8 para integración explícita, además de métodos simplécticos e implícitos. La integración se puede aplicar con vmap a todos los elementos, incluido el intervalo de integración, y el estado puede ser cualquier PyTree. También ofrece soluciones densas y varios métodos adjuntos para propagar gradientes hacia atrás a través de una resolución.
Diffrax está dirigido a quienes necesitan simulaciones diferenciables, como las ecuaciones diferenciales neuronales. Se instala con pip y requiere Python 3.10 o una versión posterior.
Funciones clave
- 01Resuelve ODE, SDE y CDE con una estructura unificada.
- 02Incluye solvers explícitos, simplécticos e implícitos, como Tsit5 y Dopri8.
- 03Permite aplicar vmap a la resolución, incluido el intervalo de integración.
- 04Acepta estados PyTree y ofrece soluciones densas.
- 05Ofrece varios métodos adjuntos para propagar gradientes hacia atrás.
Casos de uso
- Entrenar modelos de ecuaciones diferenciales neuronales que necesitan gradientes a través de un solver.
- Ejecutar simulaciones por lotes de sistemas dinámicos en GPU.
- Simular procesos estocásticos con solvers de SDE.
Inicio rápido
- 1
Instala la biblioteca con pip
pip install diffrax - 2
Define un término ODE y un solver, y llama a diffeqsolve (consulta el ejemplo rápido del README)
Antes de usarlo
- Requiere Python 3.10 o posterior.
- Está construido sobre JAX, por lo que necesitas tener JAX configurado correctamente para usarlo.
En colecciones
Temas
Preguntas frecuentes
- ¿Qué es diffrax?
- Solvers numéricos de ecuaciones diferenciales para JAX, con soporte para ODE, SDE y CDE, autodiferenciación y ejecución en GPU. Diffrax es una biblioteca de solvers numéricos de ecuaciones diferenciales basada en JAX. Resuelve ecuaciones diferenciales ordinarias (ODE), ecuaciones diferenciales estocásticas (SDE) y ecuaciones diferenciales controladas (CDE) mediante una estructura interna unificada, de modo que las tres se resuelven de la misma manera.
- ¿diffrax es gratuito? ¿Qué licencia tiene?
- diffrax es open source bajo la licencia Apache-2.0. Revisa los términos de la licencia antes de usarlo en productos comerciales.
- ¿diffrax sigue mantenido?
- Actualizado en los últimos 6 meses. El último commit de diffrax fue el 6 sept 2026.
- ¿Para qué puedo usar diffrax?
- Entrenar modelos de ecuaciones diferenciales neuronales que necesitan gradientes a través de un solver. Ejecutar simulaciones por lotes de sistemas dinámicos en GPU. Simular procesos estocásticos con solvers de SDE.
- ¿Cómo se instala diffrax?
- Instala la biblioteca con pip: pip install diffrax Define un término ODE y un solver, y llama a diffeqsolve (consulta el ejemplo rápido del README)
- ¿Dónde consigo diffrax?
- El código fuente y las instrucciones de instalación están en GitHub: https://github.com/patrick-kidger/diffrax
Resumen generado por IA a partir del README del proyecto. Las estadísticas de GitHub se actualizan a diario.
Proyectos relacionados
Compilador y runtime basados en MLIR que traducen modelos de ML a una IR unificada para desplegarlos desde centros de datos hasta dispositivos móviles y edge.
Keras 3 es un framework de deep learning que funciona con JAX, TensorFlow, PyTorch u OpenVINO para inferencia.
Una biblioteca de Python que convierte código de machine learning entre frameworks como PyTorch, TensorFlow, JAX y NumPy mediante ivy.transpile.
Una biblioteca de redes neuronales para JAX centrada en la flexibilidad, con la API Flax NNX para crear, inspeccionar y depurar modelos.
Biblioteca de modelos preentrenados para Keras 3, con modelos de texto, imagen y audio compatibles con JAX, TensorFlow y PyTorch.