{
  "abstract": "Solving large ensembles of small, independent ordinary differential equations (ODEs) is a common task in computational science and engineering, arising for example in Bayesian parameter estimation, Monte Carlo uncertainty quantification, and the integration of uncoupled physical systems such as the linearised Einstein-Boltzmann equations of cosmology. Modern graphics processing units (GPUs) with thousands of cores are well suited to such parallel workloads, but realising this potential requires careful attention to GPU-specific architectural constraints. We present modax, a Python library of GPU-accelerated ODE solvers specifically optimised for solving large ensembles of low-dimensional, independent problems. Our solvers are compatible with the JAX ecosystem but are written in a low-level, thread-based programming paradigm better suited to ensembles of highly divergent trajectories than their pure-JAX counterparts. Benchmarks on the Lorenz, Robertson and van der Pol lattice systems compare modax against diffrax, DiffEqGPU$.$jl and torchdiffeq across ensemble size, dimensionality, and trajectory divergence. Our explicit Tsit5 solver achieves 30x speed-up over diffrax on Lorenz systems, and our implicit Rodas5P solver achieves over 700x speed-up on Robertson systems, demonstrating the superiority of linearly implicit Rosenbrock-Wanner methods on GPUs. Our implicit solver scales well to higher dimensions thanks to the use of sparsity-exploiting techniques and data structures. We conclude by applying modax to three examples from the field of cosmology - Bayesian parameter estimation from primordial abundances, Monte Carlo uncertainty quantification of the global 21cm signal, and the computation of uncoupled Fourier modes in the primordial power spectrum - and demonstrate significant efficiency improvements in our ability to model physical phenomena described by low-dimensional (<200D) ODEs.",
  "abstract_shown": true,
  "affiliations": [],
  "article_type": null,
  "arxiv": "2609.35661",
  "author_total": 4,
  "authors": [
    {
      "name": "Lawrence Berry"
    },
    {
      "name": "Will Handley"
    },
    {
      "name": "Oliver Hahn"
    },
    {
      "name": "Nils Schöneberg"
    }
  ],
  "bibtex": "https://research.draconitic.com/paper/0452p3rfsg77q.bib",
  "cited_by": 0,
  "doi": "10.48550/arxiv.2609.35661",
  "fulltext_url": null,
  "id": "0452p3rfsg77q",
  "issns": [],
  "landing_url": "https://arxiv.org/abs/2609.35661",
  "language": null,
  "licence": "LicenseRef-arXiv-nonexclusive-1.0",
  "licence_label": "arXiv licence",
  "licence_source": "article",
  "licence_url": "http://arxiv.org/licenses/nonexclusive-distrib/1.0/",
  "metadata_licence": "CC0-1.0",
  "native_subjects": [
    [
      "arxiv",
      "astro-ph.CO"
    ],
    [
      "arxiv",
      "astro-ph.IM"
    ],
    [
      "arxiv",
      "physics.comp-ph"
    ],
    [
      "openalex",
      "1708|Parallel Computing and Optimization Techniques"
    ]
  ],
  "openalex": "W7214759069",
  "pdf_url": "https://arxiv.org/pdf/2609.35661",
  "pmcid": null,
  "pmid": null,
  "published": "2026-09-28",
  "publisher": null,
  "ris": "https://research.draconitic.com/paper/0452p3rfsg77q.ris",
  "sources": [
    {
      "first_seen": "2026-10-30T12:25:12Z",
      "native_id": "2609.35661",
      "source": "arxiv",
      "url": "https://arxiv.org/abs/2609.35661"
    }
  ],
  "subjects": [
    {
      "field": "astronomy",
      "origin": "arxiv:astro-ph.CO",
      "subfield": "cosmology"
    },
    {
      "field": "astronomy",
      "origin": "arxiv:astro-ph.IM",
      "subfield": "instrumentation-methods"
    },
    {
      "field": "physics",
      "origin": "arxiv:physics.comp-ph",
      "subfield": "computational"
    }
  ],
  "title": "Accelerating massive ensembles of ordinary differential equations",
  "title_text": "Accelerating massive ensembles of ordinary differential equations",
  "type": "preprint",
  "updated_at": "2026-10-30T12:25:12Z",
  "url": "https://research.draconitic.com/paper/0452p3rfsg77q",
  "venue": "arXiv",
  "versions": []
}