Skip to content

fix: allow specifying type of buffers inside MTKParameters #3585

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 38 commits into
base: master
Choose a base branch
from
Open

fix: allow specifying type of buffers inside MTKParameters #3585

wants to merge 38 commits into from
+582 −218

Conversation

AayushSabharwal
Copy link
Member

Checklist

  • Appropriate tests were added
  • Any code changes were done in a way that does not break public API
  • All documentation related to code changes were updated
  • The new code follows the
    contributor guidelines, in particular the SciML Style Guide and
    COLPRAC.
  • Any new documentation only uses public API

Additional context

Add any other context about the problem here.

ChrisRackauckas
ChrisRackauckas reviewed Apr 29, 2025
View reviewed changes
src/systems/parameter_buffer.jl Outdated
t0 = nothing, substitution_limit = 1000, floatT = nothing,
container_type = Vector)
if !(container_type <: AbstractArray)
container_type = Array
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

error?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Resolved

ChrisRackauckas added a commit to SciML/DiffEqGPU.jl that referenced this pull request May 4, 2025
@ChrisRackauckas
Add a tutorial which showcases ModelingToolkit and SII
479de67 
Fixes #330.

Currently the MWE works:

```julia
using DiffEqGPU
using OrdinaryDiffEqTsit5, ModelingToolkit, StaticArrays
using ModelingToolkit: t_nounits as t, D_nounits as D

@parameters σ ρ β
@variables x(t) y(t) z(t)

eqs = [D(D(x)) ~ σ * (y - x),
    D(y) ~ x * (ρ - z) - y,
    D(z) ~ x * y - β * z]

@mtkbuild sys = ODESystem(eqs, t) split=false

u0 = SA[D(x) => 2f0,
    x => 1f0,
    y => 0f0,
    z => 0f0]

p = SA[σ => 28f0,
    ρ => 10f0,
    β => 8f0 / 3f0]

tspan = (0f0, 100f0)
prob = ODEProblem{false}(sys, u0, tspan, p, split=true)
prob = remake(prob, p = p = SVector{10, Float32}(prob.p...))
sol = solve(prob, Tsit5())

using SymbolicIndexingInterface
p_setter = setp_oop(sys, [σ, ρ, β])

using DiffEqGPU, CUDA
function prob_func2(prob, i, repeat)
    remake(prob, p = p_setter(prob,@svector(rand(Float32,3))))
end
monteprob = EnsembleProblem(prob, prob_func = prob_func2, safetycopy = false)
sol = solve(monteprob, GPUTsit5(), EnsembleGPUKernel(CUDA.CUDABackend()),
    trajectories = 10_000)
```

But you need to `#prob = get_updated_symbolic_problem(_get_root_indp(prob), prob; kwargs...)` in DiffEqBase.

What's in here drops the `split=false` part.

We need to fix `get_updated_symbolic_problem` to not promote to `Float64` and fix static array outputs in `split=true`, i.e. SciML/ModelingToolkit.jl#3585, in order to finish this tutorial.
@ChrisRackauckas ChrisRackauckas mentioned this pull request May 4, 2025
Open
@AayushSabharwal
Copy link
Member Author

Needs SciML/DiffEqBase.jl#1151 for CI to run

@AayushSabharwal AayushSabharwal force-pushed the as/mtkparams-ptype branch 2 times, most recently from d8f1e0c to 97a1618 Compare May 12, 2025 09:34
ChrisRackauckas
ChrisRackauckas reviewed May 12, 2025
View reviewed changes
@AayushSabharwal
Copy link
Member Author

Zygote error. Love to see those.

@AayushSabharwal
Copy link
Member Author

https://github.com/SciML/ModelingToolkit.jl/actions/runs/14978208478/job/42075780307?pr=3585#step:6:1437

This says Zygote has a problem with https://github.com/SciML/SciMLBase.jl/blob/master/src/initialization.jl#L298

Which is weird because initializeprobmap is just getu which is out-of-place and worked fine before.

@ChrisRackauckas
Copy link
Member

Are you sure this PR did this? Looks like something the stuff @DhairyaLGandhi has been working on could do, since it's now a NamedTuple instead of a solution object.

@AayushSabharwal
Copy link
Member Author

AayushSabharwal commented May 12, 2025 •
edited
Loading

I don't yet see how this PR could have done it, just assuming it did because this is the first time I'm seeing the failure.

@ChrisRackauckas
Copy link
Member

I guess it is passing here: https://github.com/SciML/ModelingToolkit.jl/actions/runs/14977709554/job/42074112061?pr=3626

@DhairyaLGandhi
Copy link
Member

CI is not loading SciMLSensitivity.jl, right? It is clearly the problem type adjoint in the stacktrace.

@AayushSabharwal
Copy link
Member Author

I'm not sure what you mean by "problem type adjoint"

@DhairyaLGandhi
Copy link
Member 

 +(::@NamedTuple{f::Nothing, u0::Nothing, p::@NamedTuple{tunable::Vector{Float64}, initials::Nothing, discrete::Nothing, constant::Nothing, nonnumeric::Nothing, caches::Nothing}, problem_type::Nothing, kwargs::Nothing}, ::Base.RefValue{Any})

The named Tuple has the structure that of an ODEProblem

@AayushSabharwal
Copy link
Member Author

Right, but what does the error itself mean? initializeprobmap takes in the problem, so it makes sense that it's having trouble accumulating gradients against it. Where does the Ref{Any} come from?

@DhairyaLGandhi
Copy link
Member

It can come from a couple places, but most likely is a buffer for a mutable struct

@ChrisRackauckas
Copy link
Member

CI is not loading SciMLSensitivity.jl, right? It is clearly the problem type adjoint in the stacktrace.

It is https://github.com/SciML/ModelingToolkitNeuralNets.jl/blob/main/test/lotka_volterra.jl#L15

Could this be something your SciML/SciMLSensitivity.jl#1196 PR handles?

@DhairyaLGandhi
Copy link
Member

DhairyaLGandhi commented May 12, 2025 •
edited
Loading

No, that is specifically for the adjoint of sol.u and sol isa AbstractTimeseriesSolution (so wont get triggered for prob)

AayushSabharwal and others added 25 commits May 20, 2025 17:37
@AayushSabharwal
feat: propagate p_constructor to InitializationProblem
35a7659
@AayushSabharwal
fix: propagate iip of parent problem to InitializationProblem
fd869e1
@AayushSabharwal
fix: retain type of buffers when promoting u0/p of initialization…
b69d79e 
… problem
@AayushSabharwal
fix: handle immutable buffers in initialization
aff87fa
@AayushSabharwal
fix: handle u0_constructor, p_constructor in `remake_initializati…
ad0347b 
…on_data`
@AayushSabharwal
fix: handle immutable MTKParameters in symbolic `late_binding_update_…
926fced 
…u0_p`
@AayushSabharwal
fix: handle edge case in floating point type promotion
52ce641
@AayushSabharwal
fix: call u0_constructor on resid_prototype
57e4778
@AayushSabharwal
test: test initialization on static array problems
ba2bfff
@AayushSabharwal
build: bump SciMLBase, StochasticDelayDiffEq compat
d3ce046
@AayushSabharwal
fix: handle empty syms in concrete_getu
d351267
@AayushSabharwal
fix: fix is_update_oop passed as type
329b7d5
@AayushSabharwal
fix: fix call to remake_initialization_data
6c57557
@ChrisRackauckas @AayushSabharwal
Update Project.toml
f1e422c
@AayushSabharwal
fix: fix update_initializeprob!
7600f5f
@AayushSabharwal
fix: add f field to MockIntegrator
4631a1e
@AayushSabharwal
fix: handle discretes properly in get_mtkparameters_reconstructor
a35ae7d
@AayushSabharwal
refactor: move ChainRulesCoreExt into main package
1de4cd5
@AayushSabharwal
fix: use @ignore_derivatives inside update_initializeprob!
1dac3e7
@ChrisRackauckas @AayushSabharwal
Update Project.toml
179f0bd
@AayushSabharwal
fix: add struct IntervalNonlinearFunctionExpr
2569966
@AayushSabharwal
ci: add SciMLSensitivity/Core8 to downstream CI
0987a9f
@AayushSabharwal
build: bump SciMLBase compat
37f5320
@AayushSabharwal
fix: only promote tunables/initials in promote_u0_p if non-empty
22151e7
@AayushSabharwal
refactor: reorder tests
2c0976e
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ChrisRackauckas ChrisRackauckas ChrisRackauckas approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@AayushSabharwal @ChrisRackauckas @DhairyaLGandhi