Update to using SII and SciMLStructures #1057

DhairyaLGandhi · 2024-05-22T08:19:05Z

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

Moves SciMLSensitivity to use the SII and SciMLStructures' interface.

Add any other context about the problem here.

This uses the SciMLStructures Tunable interface https://github.com/SciML/SciMLStructures.jl in order to allow more generalized definitions of `p`. - [ ] Ensure Lux.jl is well supported (componentarrays extension in SciMLStructures - [ ] Add a test for a custom SciMLStructure

AayushSabharwal

Just minor things

src/adjoint_common.jl

AayushSabharwal · 2024-05-23T15:50:59Z

src/forward_sensitivity.jl

@@ -158,7 +158,7 @@ function ODEForwardSensitivityProblem(f::F, args...; kwargs...) where {F}
 end

 function ODEForwardSensitivityProblem(prob::ODEProblem, alg; kwargs...)
- ODEForwardSensitivityProblem(prob.f, prob.u0, prob.tspan, prob.p, alg; kwargs...)
+ ODEForwardSensitivityProblem(symbolic_container(prob), state_values(prob), prob.tspan, parameter_values(prob), alg; kwargs...)


symbolic_container technically does returns prob.f, however semantically it isn't guaranteed to do so and in such cases where you explicitly wand the function, I'd still use prob.f.

This has been addressed

Co-authored-by: Aayush Sabharwal <[email protected]>

…into dg/structures

ChrisRackauckas · 2024-05-26T00:03:48Z

src/concrete_solve.jl

@@ -354,14 +361,15 @@ function DiffEqBase._concrete_solve_adjoint(
 sol = solve(_prob, alg, args...; save_noise = true, save_start = true,
 save_end = true, kwargs_fwd...)
 end
+ time = current_time(sol)


what's current time from?

SII https://github.com/SciML/SymbolicIndexingInterface.jl/blob/71408bb7f6f22791144542061f05287f65e9b700/src/value_provider_interface.jl#L121

ChrisRackauckas · 2024-05-26T00:04:13Z

src/concrete_solve.jl

- ts = sol.t
+ ts = time


Why change the standard naming? Naming an array time instead of a scalar value is weird and not grammatically correct...

I used it where there were multiple calls to sol.t. I wanted to avoid multiple calls to current_time when something like sol.t was referenced multiple times, since field access might be free, function calls potentially may not be.

I can change the naming or potentially if there's a better solution, can implement that

That's fine, it could still be called ts. If you want to write it, times. But I don't understand why syntactic changes are mixed with a feature PR.

I can change it to be ts - being shorter to write.

This should be fixed now with 6952e1b and 95cad85

Project.toml

…into dg/structures

…ners

…into dg/structures

DhairyaLGandhi · 2024-06-25T21:07:26Z

The shadowing tests are passing now as well. The MWE for the DiffEqFlux failures is:

julia> ca = ComponentVector((x = rand(3), y = rand(4),))  # start with a component array                                                                                                          
ComponentVector{Float64}(x = [0.2755239169592659, 0.2144463395930084, 0.9085629964224934], y = [0.122401457810142, 0.498989539972836, 0.7718946152650026, 0.9715871491699061]
)                                                                                                                                                                            
                                                                                                                                                                             
julia> tr = Tracker.param(ca)  # use Tracker as the backend                                                                                                                                   
ComponentVector{Float64, TrackedVector{Float64, Vector{Float64}}, Tuple{Axis{(x = 1:3, y = 4:7)}}}(x = [0.2755239169592659, 0.2144463395930084, 0.9085629964224934] (tracked), y = [0.122401457810142, 0.498989539972836, 0.7718946152650026, 0.9715871491699061] (tracked))

julia> convert(typeof(ca), tr) # convert call masks the type of the underlying array
Tracked 7-element ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(x = 1:3, y = 4:7)}}} with indices 1:1:7:
 0.2755239169592659
 0.2144463395930084
 0.9085629964224934
 0.122401457810142
 0.498989539972836
 0.7718946152650026
 0.9715871491699061

julia> Float64.(tr)
Tracked 7-element ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(x = 1:3, y = 4:7)}}} with indices 1:1:7:
 0.2755239169592659
 0.2144463395930084
 0.9085629964224934
 0.122401457810142
 0.498989539972836
 0.7718946152650026
 0.9715871491699061

The specific convert method is https://github.com/jonniedie/ComponentArrays.jl/blob/554a9c03373680af84586762f68ebd32b6d34abe/src/similar_convert_copy.jl#L55 which seems to be missing a case where T1 === T2 and does an extra allocation. A dispatch could be defined to avoid this.

function Base.convert(::Type{ComponentArray{T,N,A1,Ax1}}, x::ComponentArray{T,N,A2,Ax2}) where {T,N,A1,A2,Ax1,Ax2}
    return x
end

ChrisRackauckas · 2024-06-26T00:57:24Z

Can you PR that? If @jonniedie has a quick turnaround then I think we're close to done here.

DhairyaLGandhi · 2024-06-26T10:15:10Z

Ref jonniedie/ComponentArrays.jl#263

DhairyaLGandhi · 2024-06-26T12:53:48Z

src/gauss_adjoint.jl


 if !isautojacvec
 if DiffEqBase.has_paramjac(f)
 f.paramjac(pJ, y, p, t) # Calculate the parameter Jacobian into pJ
 else
 pf.t = t
+ pf.u = y


While working through this PR, I found this path which was not updating the caches. I have added this here, but wanted to highlight for correctness

ChrisRackauckas · 2024-06-27T20:53:07Z

@wsmoses we're seeing something odd with the latest Enzyme release as well.

wsmoses · 2024-06-27T20:55:46Z

@ChrisRackauckas how so, what’s the failure ?

ChrisRackauckas · 2024-06-27T20:58:31Z

same as @avik-pal is mentioning is seen in CI here

Continuous and discrete costs: Error During Test at /home/runner/.julia/packages/SafeTestsets/raUNr/src/SafeTestsets.jl:30
  Got exception outside of a @test
  LoadError: MethodError: no method matching cpu_features!(::LLVM.ModulePassManager)
  
  Closest candidates are:
    cpu_features!(::LLVM.Module)
     @ GPUCompiler ~/.julia/packages/GPUCompiler/05oYT/src/optim.jl:606
  
  Stacktrace:
    [1] (::Enzyme.Compiler.var"#28472#28473"{LLVM.Module, LLVM.TargetMachine})(pm::LLVM.ModulePassManager)
      @ Enzyme.Compiler ~/.julia/packages/Enzyme/qd8AI/src/compiler/optimize.jl:1941
    [2] LLVM.ModulePassManager(::Enzyme.Compiler.var"#28472#28473"{LLVM.Module, LLVM.TargetMachine}; kwargs::@Kwargs{})
      @ LLVM ~/.julia/packages/LLVM/5DlHM/src/passmanager.jl:33
    [3] LLVM.ModulePassManager(::Function)
      @ LLVM ~/.julia/packages/LLVM/5DlHM/src/passmanager.jl:30
    [4] optimize!(mod::LLVM.Module, tm::LLVM.TargetMachine)

wsmoses · 2024-06-27T20:59:52Z

Hm odd, post that backtrace as an issue and I’ll try to solve it tomorrow

…

On Thu, Jun 27, 2024 at 10:58 PM Christopher Rackauckas < ***@***.***> wrote: same as @avik-pal <https://github.com/avik-pal> is mentioning is seen in CI here Continuous and discrete costs: Error During Test at /home/runner/.julia/packages/SafeTestsets/raUNr/src/SafeTestsets.jl:30 Got exception outside of a @test LoadError: MethodError: no method matching cpu_features!(::LLVM.ModulePassManager) Closest candidates are: cpu_features!(::LLVM.Module) @ GPUCompiler ~/.julia/packages/GPUCompiler/05oYT/src/optim.jl:606 Stacktrace: [1] (::Enzyme.Compiler.var"#28472#28473"{LLVM.Module, LLVM.TargetMachine})(pm::LLVM.ModulePassManager) @ Enzyme.Compiler ~/.julia/packages/Enzyme/qd8AI/src/compiler/optimize.jl:1941 [2] LLVM.ModulePassManager(::Enzyme.Compiler.var"#28472#28473"{LLVM.Module, LLVM.TargetMachine}; ***@***.***{}) @ LLVM ~/.julia/packages/LLVM/5DlHM/src/passmanager.jl:33 [3] LLVM.ModulePassManager(::Function) @ LLVM ~/.julia/packages/LLVM/5DlHM/src/passmanager.jl:30 [4] optimize!(mod::LLVM.Module, tm::LLVM.TargetMachine) — Reply to this email directly, view it on GitHub <#1057 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJTUXCEIZE6D3KWULVVK3TZJR4I3AVCNFSM6AAAAABIDDSPO6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCOJVGY2TKOJWHA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

wsmoses · 2024-06-27T21:01:17Z

Oh I think this is the llvm.jl bump changing api (it was done properly with a change which did semver but it passed ci so I assume we didn’t hit).

…

On Thu, Jun 27, 2024 at 10:59 PM Billy Moses ***@***.***> wrote: Hm odd, post that backtrace as an issue and I’ll try to solve it tomorrow On Thu, Jun 27, 2024 at 10:58 PM Christopher Rackauckas < ***@***.***> wrote: > same as @avik-pal <https://github.com/avik-pal> is mentioning is seen in > CI here > > Continuous and discrete costs: Error During Test at /home/runner/.julia/packages/SafeTestsets/raUNr/src/SafeTestsets.jl:30 > Got exception outside of a @test > LoadError: MethodError: no method matching cpu_features!(::LLVM.ModulePassManager) > > Closest candidates are: > cpu_features!(::LLVM.Module) > @ GPUCompiler ~/.julia/packages/GPUCompiler/05oYT/src/optim.jl:606 > > Stacktrace: > [1] (::Enzyme.Compiler.var"#28472#28473"{LLVM.Module, LLVM.TargetMachine})(pm::LLVM.ModulePassManager) > @ Enzyme.Compiler ~/.julia/packages/Enzyme/qd8AI/src/compiler/optimize.jl:1941 > [2] LLVM.ModulePassManager(::Enzyme.Compiler.var"#28472#28473"{LLVM.Module, LLVM.TargetMachine}; ***@***.***{}) > @ LLVM ~/.julia/packages/LLVM/5DlHM/src/passmanager.jl:33 > [3] LLVM.ModulePassManager(::Function) > @ LLVM ~/.julia/packages/LLVM/5DlHM/src/passmanager.jl:30 > [4] optimize!(mod::LLVM.Module, tm::LLVM.TargetMachine) > > — > Reply to this email directly, view it on GitHub > <#1057 (comment)>, > or unsubscribe > <https://github.com/notifications/unsubscribe-auth/AAJTUXCEIZE6D3KWULVVK3TZJR4I3AVCNFSM6AAAAABIDDSPO6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCOJVGY2TKOJWHA> > . > You are receiving this because you were mentioned.Message ID: > ***@***.***> >

avik-pal · 2024-06-27T21:04:49Z

Oh I think this is the llvm.jl bump changing api (it was done properly with a change which did semver but it passed ci so I assume we didn’t hit).

On enzyme CI it installed v7 so didn't really check the correct thing.

Really wish there was a way to remove old deps in CompatHelper PRs and push another change with the old deps inplace once tests pass. (I would have guessed the force_latest_versions kwarg in test does this)

wsmoses · 2024-06-27T21:09:43Z

@ChrisRackauckas @avik-pal does this fix?

EnzymeAD/Enzyme.jl#1581

ChrisRackauckas · 2024-06-27T21:59:07Z

src/adjoint_common.jl

@@ -34,12 +34,14 @@ function adjointdiffcache(g::G, sensealg, discrete, sol, dgdu::DG1, dgdp::DG2, f
 prob = sol.prob
 u0 = state_values(prob)
 p = parameter_values(prob)
- if p === nothing || p isa SciMLBase.NullParameters
+ if p === nothing || p === DiffEqBase.NullParameters()


why the change to DiffEq? Should be SciML

ChrisRackauckas and others added 2 commits February 8, 2024 11:15

chore: update to SciMLStructuers

5abcec3

DhairyaLGandhi changed the title ~~Dg/structures~~ Update to using SII and SciMLStructures May 22, 2024

DhairyaLGandhi added 9 commits May 22, 2024 17:39

chore: update gauss adjoint to SII

4a3690a

chore: update adjointproblem to SII

9f69618

chore: use current_time

5c4cf5e

chore: add current_time to concrete_solve

5ef9e89

chore: crlf

adb909e

chore: line endings

f560342

chore: add SII, SStructures to deps

795af73

chore: merge conflicts

ac27e3f

chore: fix lux compat

9fe8f50

AayushSabharwal reviewed May 23, 2024

View reviewed changes

DhairyaLGandhi and others added 4 commits May 23, 2024 21:25

Update src/adjoint_common.jl

e727867

Co-authored-by: Aayush Sabharwal <[email protected]>

chore: access .f directly

59415c1

Merge branch 'dg/structures' of github.com:SciML/SciMLSensitivity.jl …

b9ce342

…into dg/structures

chore: typo

5e733a1

ChrisRackauckas reviewed May 26, 2024

View reviewed changes

Project.toml Outdated Show resolved Hide resolved

ChrisRackauckas and others added 10 commits May 26, 2024 09:56

Update Project.toml

e5410ec

chore: use correct dispatch for automatic_sensealg_choice

115ec2a

chore: use corrent inplace_vjp dispatch

a16aa98

chore: consistent names

e6b4fe6

Merge branch 'dg/structures' of github.com:SciML/SciMLSensitivity.jl …

f12bc2a

…into dg/structures

chore: fix imports

69f217d

chore: pick u0 correctly

df65237

test: use SciMLStructures interface in tests

068c133

chore: use SII for sensitivity_solution

fea9e7c

test: use correct inplace_vjp dipatch

d1c3cec

DhairyaLGandhi added 9 commits June 24, 2024 19:49

chore: move enzyme error message back

b6db6f8

chore: use sol.prob to access u0

4ef4a9e

chore: use ArrayInterface.aos_to_soa to handle mutable tracked contai…

6adc3f1

…ners

Merge branch 'dg/structures' of github.com:SciML/SciMLSensitivity.jl …

1b549c7

…into dg/structures

chore: revert to using views for accessing state

c95633c

chore: typo

939634f

chore: clean up shadowing tests

6239ae1

chore: format

2577176

chore: cleanup

ba214c3

chore: replace rather than repack

783cc82

DhairyaLGandhi commented Jun 27, 2024

View reviewed changes

DhairyaLGandhi added 2 commits June 28, 2024 00:46

chore: move error message back

26d1465

chore: use ArrayInterface for g

3559012

ChrisRackauckas mentioned this pull request Jun 27, 2024

LLVM version issue with latest release? EnzymeAD/Enzyme.jl#1580

Closed

DhairyaLGandhi added 2 commits June 28, 2024 03:14

chore: SciMLBase.NullParameters -> DiffEqBase.NullParameters

a752804

chore: protect against wrong p in adjointdiffcache

9f888fe

ChrisRackauckas reviewed Jun 27, 2024

View reviewed changes

chore: revert DiffEqBase.NullParameters

0ce3753

ChrisRackauckas merged commit 903f96d into master Jun 28, 2024
15 of 16 checks passed

ChrisRackauckas deleted the dg/structures branch June 28, 2024 00:14

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update to using SII and SciMLStructures #1057

Update to using SII and SciMLStructures #1057

DhairyaLGandhi commented May 22, 2024

AayushSabharwal left a comment

AayushSabharwal May 23, 2024

DhairyaLGandhi May 27, 2024

ChrisRackauckas May 26, 2024

DhairyaLGandhi May 26, 2024

ChrisRackauckas May 26, 2024

DhairyaLGandhi May 26, 2024

DhairyaLGandhi May 26, 2024

ChrisRackauckas May 26, 2024

DhairyaLGandhi May 26, 2024

DhairyaLGandhi May 28, 2024

DhairyaLGandhi commented Jun 25, 2024 •

edited

Loading

ChrisRackauckas commented Jun 26, 2024

DhairyaLGandhi commented Jun 26, 2024

DhairyaLGandhi Jun 26, 2024

ChrisRackauckas commented Jun 27, 2024

wsmoses commented Jun 27, 2024

ChrisRackauckas commented Jun 27, 2024

wsmoses commented Jun 27, 2024 via email

wsmoses commented Jun 27, 2024 via email

avik-pal commented Jun 27, 2024 •

edited

Loading

wsmoses commented Jun 27, 2024

ChrisRackauckas Jun 27, 2024

DhairyaLGandhi Jun 27, 2024

Update to using SII and SciMLStructures #1057

Update to using SII and SciMLStructures #1057

Conversation

DhairyaLGandhi commented May 22, 2024

Checklist

Additional context

AayushSabharwal left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DhairyaLGandhi commented Jun 25, 2024 • edited Loading

ChrisRackauckas commented Jun 26, 2024

DhairyaLGandhi commented Jun 26, 2024

Choose a reason for hiding this comment

ChrisRackauckas commented Jun 27, 2024

wsmoses commented Jun 27, 2024

ChrisRackauckas commented Jun 27, 2024

wsmoses commented Jun 27, 2024 via email

wsmoses commented Jun 27, 2024 via email

avik-pal commented Jun 27, 2024 • edited Loading

wsmoses commented Jun 27, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DhairyaLGandhi commented Jun 25, 2024 •

edited

Loading

avik-pal commented Jun 27, 2024 •

edited

Loading