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
Merged
merged 40 commits into from
May 21, 2025
Merged

fix: allow specifying type of buffers inside MTKParameters #3585

merged 40 commits into from
May 21, 2025

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
Merged
@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 16 commits May 20, 2025 17:37
@AayushSabharwal
Copy link
Member Author

MTKNN and SciMLSensitivity failures are the same as the one in Extensions

@ChrisRackauckas ChrisRackauckas merged commit d91b8d1 into SciML:master May 21, 2025
31 of 47 checks passed
@AayushSabharwal AayushSabharwal deleted the as/mtkparams-ptype branch May 21, 2025 10:24
ChrisRackauckas added a commit to SciML/DiffEqGPU.jl that referenced this pull request May 21, 2025
@ChrisRackauckas
Add a tutorial which showcases ModelingToolkit and SII
d2db6f5 
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.
@DhairyaLGandhi DhairyaLGandhi mentioned this pull request May 25, 2025
Merged
2 tasks
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