Update ensemble interface for SciMLBase v3#432
Merged
ChrisRackauckas merged 1 commit intoSciML:masterfrom Apr 12, 2026
Merged
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
SciMLBasecompat bound from2.144to3Breaking Changes Addressed
SciMLBase v3 changed:
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
EnsembleContextbundlessim_id,repeat,worker_id,sim_seed,rng, andmaster_rng.Changes
src/solve.jl:__solve: acceptsseed,rng,rng_funckwargs; generates per-trajectory seeds viaSciMLBase.generate_sim_seeds; bundlesensemble_rng_statefor CPU offload pathbatch_solve: createsEnsembleContextvia new_make_ensemble_contexthelper for each trajectory; callsprob_func(prob, ctx)andoutput_func(sol, ctx)with context instead of raw indicessolve_batch(EnsembleThreads override): accepts and forwardsensemble_rng_statepositional arg to match SciMLBase v3'ssolve_batchsignaturesrc/algorithms.jl: Updated docstring examplestest/: Updated all 17 test files to use new(prob, ctx)/(sol, ctx)signatures, replacingiwithctx.sim_idwhere the trajectory index was usedProject.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_solvethat bypass SciMLBase's defaultbatch_func, so the newEnsembleContextcreation 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
prob_funcprob_func(prob, ctx)withctx.sim_idworksseedkwarg produces reproducible GPU ensemble resultsensemble_rng_statecorrectly🤖 Generated with Claude Code