Skip to content

Update ensemble interface for SciMLBase v3#432

Merged
ChrisRackauckas merged 1 commit intoSciML:masterfrom
ChrisRackauckas-Claude:update-scimlbase-v3-ensemble-interface
Apr 12, 2026
Merged

Update ensemble interface for SciMLBase v3#432
ChrisRackauckas merged 1 commit intoSciML:masterfrom
ChrisRackauckas-Claude:update-scimlbase-v3-ensemble-interface

Conversation

@ChrisRackauckas-Claude
Copy link
Copy Markdown
Contributor

Summary

  • Updates DiffEqGPU's ensemble solve implementation to match SciMLBase v3's new ensemble interface
  • SciMLBase v3 (via SciML/SciMLBase.jl#1252) redesigned the ensemble RNG system with breaking signature changes
  • Updates SciMLBase compat bound from 2.144 to 3

Breaking Changes Addressed

SciMLBase v3 changed:

v2 (old) v3 (new)
prob_func (prob, i, repeat) (prob, ctx::EnsembleContext)
output_func (sol, i) (sol, ctx::EnsembleContext)
solve_batch (prob, alg, ensemblealg, II, pmap_batch_size; kwargs...) (prob, alg, ensemblealg, II, pmap_batch_size, ensemble_rng_state; kwargs...)

Where EnsembleContext bundles sim_id, repeat, worker_id, sim_seed, rng, and master_rng.

Changes

src/solve.jl:

  • __solve: accepts seed, rng, rng_func kwargs; generates per-trajectory seeds via SciMLBase.generate_sim_seeds; bundles ensemble_rng_state for CPU offload path
  • batch_solve: creates EnsembleContext via new _make_ensemble_context helper for each trajectory; calls prob_func(prob, ctx) and output_func(sol, ctx) with context instead of raw indices
  • solve_batch (EnsembleThreads override): accepts and forwards ensemble_rng_state positional arg to match SciMLBase v3's solve_batch signature

src/algorithms.jl: Updated docstring examples

test/: Updated all 17 test files to use new (prob, ctx) / (sol, ctx) signatures, replacing i with ctx.sim_id where the trajectory index was used

Project.toml: SciMLBase = "3"

Note

This PR supersedes #431 (dependabot compat-only bump). A compat bump alone is insufficient because DiffEqGPU implements its own __solve/batch_solve that bypass SciMLBase's default batch_func, so the new EnsembleContext creation and RNG infrastructure must be replicated here.

CI will not resolve until downstream deps (e.g., LinearSolve) also bump their SciMLBase compat to v3.

Test plan

  • CI passes once LinearSolve and other deps are updated for SciMLBase v3
  • Verify GPU ensemble solves produce correct results with default prob_func
  • Verify custom prob_func(prob, ctx) with ctx.sim_id works
  • Verify seed kwarg produces reproducible GPU ensemble results
  • Verify CPU offload path forwards ensemble_rng_state correctly

🤖 Generated with Claude Code

SciMLBase v3 changed the ensemble callback signatures:
- prob_func(prob, i, repeat) -> prob_func(prob, ctx::EnsembleContext)
- output_func(sol, i) -> output_func(sol, ctx::EnsembleContext)
- solve_batch now requires ensemble_rng_state argument

This updates DiffEqGPU's custom __solve/batch_solve to:
- Accept seed/rng/rng_func kwargs and generate per-trajectory seeds
- Create EnsembleContext objects and pass to prob_func/output_func
- Forward ensemble_rng_state to SciMLBase.solve_batch for CPU offload
- Update all test files to use new 2-arg signatures

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@ChrisRackauckas ChrisRackauckas merged commit 0c1915d into SciML:master Apr 12, 2026
4 of 25 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants