JAX wrapper around diffsol, a Rust ODE solver library.
Exposes diffsol's ODE solvers via jax.ffi so they can be called from inside jax.jit and
differentiated with jax.grad.
The user writes an RHS function in Python, which gets lowered to a diffsl source string and compiled on first call.
Example usage
import jax
import jax.numpy as jnp
from diffsol_jax import ODEProblem
def lotka_volterra(t, y, p):
x, yy = y[0], y[1]
alpha, beta, delta, gamma = p[0], p[1], p[2], p[3]
return (alpha * x - beta * x * yy, delta * x * yy - gamma * yy)
params = jnp.array([1.5, 1.0, 0.75, 3.0])
y0 = jnp.array([1.0, 0.5])
problem = ODEProblem(lotka_volterra, y0=y0, params=params)
t_span = jnp.array([0.0, 10.0])
ts, ys = jax.jit(lambda p: problem.solve(p, t_span))(params)
# ts: float64[200], ys: float64[200, 2]
Limitations
- CPU only
- No vmap batching rule;
vmap_method="sequential"gives correct results via a Python loop. - The DiffSL lowerer handles elementwise ops and the common ODE patterns (Lotka-Volterra, Lorenz).
Operations like
dot_general,reduce_sum, andconcatenateare not supported, so neural ODEs are not supported yet. - Differentiation uses JAX-level forward sensitivities (no separate adjoint pass). DiffSL modules
are JIT-compiled with Cranelift and cached per source string; the first call compiles, subsequent
calls reuse the compiled module and update parameters via
set_params. - Does not work with
jax.pmap - Higher-order derivatives
1 maintainer
Metadata
1 maintainer
Release files for diffsol-jax 0.0.2
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| diffsol_jax-0.0.2.tar.gz | 880.3 kB | Details |
Built distributions (wheels)
Total release size: 48.9 MB
Release files / diffsol_jax-0.0.2.tar.gz
| Download URL | diffsol_jax-0.0.2.tar.gz |
|---|---|
| Size | 880.3 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
8388eafbb77176a6e8032ef1a947d6d12a10e79ef861d6f772db72075541c950
|
|
BLAKE2b-256 checksum How to use checksums |
5c509c5dd32239ebfe9a118ad9791684ccb5c06ed196a4edbaf50253623567ae
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency logRelease files / diffsol_jax-0.0.2-cp313-cp313-win_amd64.whl
| Download URL | diffsol_jax-0.0.2-cp313-cp313-win_amd64.whl |
|---|---|
| Size | 3.7 MB |
| Tags | CPython 3.13 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
cc149f3c02c47bfc7c927cd9f0fe492f1b264eb260a2a995474bc5d8c7e3c23c
|
|
BLAKE2b-256 checksum How to use checksums |
a89b99bf7e04f10ce4570b55e232220314fde3b9bd02c71005dc9102b6559f0a
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency logRelease files / diffsol_jax-0.0.2-cp313-cp313-manylinux_2_28_x86_64.whl
| Download URL | diffsol_jax-0.0.2-cp313-cp313-manylinux_2_28_x86_64.whl |
|---|---|
| Size | 7.0 MB |
| Tags | CPython 3.13 Linux glibc 2.28+ x86-64 |
|
SHA-256 checksum How to use checksums |
f4823f549a77808d6ba5a2693ca25b7e1cf03afc95b530ffd4a0059af79c7201
|
|
BLAKE2b-256 checksum How to use checksums |
5684f27c8651839b2c98badd7defbe2986bbffe623b3184c8ee80c31b525240f
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency logRelease files / diffsol_jax-0.0.2-cp313-cp313-macosx_14_0_arm64.whl
| Download URL | diffsol_jax-0.0.2-cp313-cp313-macosx_14_0_arm64.whl |
|---|---|
| Size | 5.3 MB |
| Tags | CPython 3.13 macOS 14.0+ ARM64 |
|
SHA-256 checksum How to use checksums |
041c88420afafe9da9dbdb4557ca269a618563e715e77eebd28e34fc912d4b50
|
|
BLAKE2b-256 checksum How to use checksums |
5e5968439f446b2845ef3ce527be055cf2bf160779e2117e2cef5d05df80abe2
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency logRelease files / diffsol_jax-0.0.2-cp312-cp312-win_amd64.whl
| Download URL | diffsol_jax-0.0.2-cp312-cp312-win_amd64.whl |
|---|---|
| Size | 3.7 MB |
| Tags | CPython 3.12 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
b41bf02190f4e0d5ea59380666c5afc86fa42bbc518ad2421d106fc55a6f2dd5
|
|
BLAKE2b-256 checksum How to use checksums |
a13facf92766f451f0c612e4a2d5d05fb987b149ba24461bc04b056628d85f6f
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency logRelease files / diffsol_jax-0.0.2-cp312-cp312-manylinux_2_28_x86_64.whl
| Download URL | diffsol_jax-0.0.2-cp312-cp312-manylinux_2_28_x86_64.whl |
|---|---|
| Size | 7.0 MB |
| Tags | CPython 3.12 Linux glibc 2.28+ x86-64 |
|
SHA-256 checksum How to use checksums |
7b381040d2a8442d22422be387d166e72864f67e46e886b8e6416770832119a2
|
|
BLAKE2b-256 checksum How to use checksums |
2191f6612c58dc6e0f1d4ec9a5bc5d7a8e23dd9cc2a6e16e759deeb550e77e25
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency logRelease files / diffsol_jax-0.0.2-cp312-cp312-macosx_14_0_arm64.whl
| Download URL | diffsol_jax-0.0.2-cp312-cp312-macosx_14_0_arm64.whl |
|---|---|
| Size | 5.3 MB |
| Tags | CPython 3.12 macOS 14.0+ ARM64 |
|
SHA-256 checksum How to use checksums |
3810e1ca66c9a68786d2b49106af51a59cd631623b5c833749c766db4be0955b
|
|
BLAKE2b-256 checksum How to use checksums |
b041b31369ef4da70ffe940fdf578004492dc1924ce9aaf4f1b3b0bb1c7c91de
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency logRelease files / diffsol_jax-0.0.2-cp311-cp311-win_amd64.whl
| Download URL | diffsol_jax-0.0.2-cp311-cp311-win_amd64.whl |
|---|---|
| Size | 3.7 MB |
| Tags | CPython 3.11 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
5c9655c58de2dcf957ed0f0bcb2bf5fe735b8c796ea88d6929acbb8ce2520bf3
|
|
BLAKE2b-256 checksum How to use checksums |
0888824287e602a20f0842350935023b3626fda8ca2c1c1a7865a739408f38e5
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency logRelease files / diffsol_jax-0.0.2-cp311-cp311-manylinux_2_28_x86_64.whl
| Download URL | diffsol_jax-0.0.2-cp311-cp311-manylinux_2_28_x86_64.whl |
|---|---|
| Size | 7.0 MB |
| Tags | CPython 3.11 Linux glibc 2.28+ x86-64 |
|
SHA-256 checksum How to use checksums |
8aeaa14a0971b8baecc0868801c6495bb32f959718f3419df97dcdfdfbdeaa7b
|
|
BLAKE2b-256 checksum How to use checksums |
74dc8eaaa71c94bd1da3a169a0e581d75ddb244e287a5384d19b076feec01d4f
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency logRelease files / diffsol_jax-0.0.2-cp311-cp311-macosx_14_0_arm64.whl
| Download URL | diffsol_jax-0.0.2-cp311-cp311-macosx_14_0_arm64.whl |
|---|---|
| Size | 5.3 MB |
| Tags | CPython 3.11 macOS 14.0+ ARM64 |
|
SHA-256 checksum How to use checksums |
e29799604c939190c3fc61f77ff34ff22d03c955a374a507f737db2c97eb7f7a
|
|
BLAKE2b-256 checksum How to use checksums |
d810fcabc254468274f839dd4f2794f5833246f35fe0fc067f5f0d081b7bbf97
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Jun 24, 2026.
Transparency log
