Harmonic Oscillator Decapode Calibration

DARPA-ASKEM / sciml-service

Simulation Service provides an interface and job runner for ASKEM models.

MIT License

3 stars 1 forks source link

Harmonic Oscillator Decapode Calibration #177

Open jClugstor opened 1 month ago

jClugstor commented 1 month ago

An effort to create a very simple Decapode, to help simplify diagnosing autodiff through decapode issues. This is a 1-D harmonic oscillator Decapode:

using Catlab
using Catlab.Graphics
using CombinatorialSpaces
using Decapodes
using ComponentArrays

using Catlab
using Catlab.Graphics
using CombinatorialSpaces
using CombinatorialSpaces.ExteriorCalculus
using DiagrammaticEquations
using DiagrammaticEquations.Deca
using Decapodes
using GeometryBasics:Point1,Point2,Point3
using OrdinaryDiffEq

Point1D = Point1{Float64}
Point2D = Point2{Float64};
Point3D = Point3{Float64};

oscillator = @decapode begin
    X::Form0
    V::Form0

    k::Constant

    ∂ₜ(X) == V
    ∂ₜ(V) == -k*(X)
end

decapode_code = gensim(oscillator, [:X, :k, :V], dimension = 1)

file = open("oscillator_code.jl", "w")
write(file, string("decapode_f = ", decapode_code))
close(file)
include("oscillator_code.jl")

x0 = [1.0]
v0 = [0.0]

u₀ = ComponentArray(X = x0, V = v0)
p = (k = 1.0,)
tₑ = 10

s_prime = EmbeddedDeltaSet1D{Bool, Point1D}()
add_vertices!(s_prime,1)

s = EmbeddedDeltaDualComplex1D{Bool, Float64, Point1D}(s_prime)

function generate(operators,mesh)
end

fₘ = decapode_f(s, generate)

data_prob = ODEProblem{true, SciMLBase.FullSpecialize}(fₘ, u₀, (0, tₑ),p)
sol = solve(data_prob,Tsit5(), saveat = 0.1)
dat = vcat([u.X for u in sol.u]...)
CairoMakie.lines(dat)

using SciMLSensitivity, Zygote, Enzyme

function fake_loss(constants_and_parameters)
  prob = remake(data_prob, p = constants_and_parameters)
  soln = solve(prob, FBDF(autodiff = false))
  @info("Done")

  # return soln(tₑ)
  sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end
Zygote.gradient(fake_loss, p)

Zygote.gradient(fake_loss,p) errors with

ArgumentError: cannot reinterpret an `ForwardDiff.Dual{nothing, Float64, 1}` array to `ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.FullSpecialize, var"#f#7"{PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Float64}, Float64, 2}` whose first dimension has size `1`.
The resulting array would have non-integral first dimension.

is the ForwardDiff coming from the solver adjoint?

Anyways, setting the sensealg: soln = solve(prob, FBDF(autodiff = false),sensealg = SciMLSensitivity.ZygoteVJP) and running Zygote.gradient(fake_loss,p) gets a different error inside of Zygote: ERROR: ArgumentError: new: too few arguments (expected 49)

jClugstor commented 1 month ago

@jpfairbanks mentioned today at the PI kickoff meeting that maybe the best way forward is to use SciMLStructures to set the desired Decapode parameters to be tunable. I also tried Enzyme, did not work. @ChrisRackauckas

ChrisRackauckas commented 1 month ago

What is the error you get from Enzyme?

maybe the best way forward is to use SciMLStructures to set the desired Decapode parameters to be tunable

Yes, though note that will need to use the PR branch https://github.com/SciML/SciMLSensitivity.jl/pull/1057 until it merges (hopefully in the next week or so)

jClugstor commented 1 month ago

Enzyme.gradient(Enzyme.Reverse, fake_loss, p) gives:

Enzyme execution failed.
Mismatched activity for:   store i64 %11, i64 addrspace(11)* %10, align 8, !dbg !42 const val:   %11 = load i64, i64 addrspace(11)* %9, align 8, !dbg !42
Type tree: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

inside of the remake call.

ChrisRackauckas commented 1 month ago

Can we MWE this more to not include Decapodes?

jClugstor commented 1 month ago

This is a version that uses a function that was originally generated by Decapodes, but has any Decapode stuff and CombinatorialSpaces stuff stripped out, but gives the same results when solved. Trying to autodiff gives roughly the same errors as before.

using ComponentArrays
using PreallocationTools
using OrdinaryDiffEq

not_decapode_f = begin

    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
            end
    end
end

x0 = [1.0]
v0 = [0.0]

u₀ = ComponentArray(X = x0, V = v0)
p = (k = [1.0],)
tₑ = 10

fₘ = not_decapode_f()

data_prob = ODEProblem{true, SciMLBase.FullSpecialize}(fₘ, u₀, (0, tₑ),p)
sol = solve(data_prob,Tsit5(), saveat = 0.1)
dat = vcat([u.X for u in sol.u]...)

using SciMLSensitivity, Zygote, Enzyme

function fake_loss_default_number(p)
    newp = (k = p,)
    prob = remake(data_prob, p = newp)
    soln = solve(prob, Tsit5())
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

function fake_loss_default(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, Tsit5())
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss_default_number, [1.0])
Zygote.gradient(fake_loss_default, p)
Enzyme.gradient(Enzyme.Reverse, fake_loss_default, p)
Enzyme.gradient(Enzyme.Reverse, fake_loss_default, p)

function fake_loss_Zygote(p)
  prob = remake(data_prob, p = p)
  soln = solve(prob, FBDF(autodiff = false), sensealg = ZygoteVJP())
  @info("Done")

  # return soln(tₑ)
  sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

function fake_loss_Enzyme(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, FBDF(autodiff = false), sensealg = EnzymeVJP())
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss_Zygote, p)
Zygote.gradient(fake_loss_Enzyme, p)
Enzyme.gradient(Enzyme.Reverse, fake_loss_Zygote, p)
Enzyme.gradient(Enzyme.Reverse, fake_loss_Enzyme, p)

function fake_loss_Zygote_number_p(p)
    newp = (k = p,)
    prob = remake(data_prob, p = newp)
    soln = solve(prob, FBDF(autodiff = false), sensealg = ZygoteVJP())
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

function fake_loss_Enzyme_number_p(p)
    newp = (k = p,)
    prob = remake(data_prob, p = newp)
    soln = solve(prob, FBDF(autodiff = false), sensealg = EnzymeVJP())
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss_Zygote_number_p, 1.0)
Zygote.gradient(fake_loss_Enzyme_number_p, 1.0)
Enzyme.gradient(Enzyme.Reverse, fake_loss_Zygote_number_p, 1.0)
Enzyme.gradient(Enzyme.Reverse, fake_loss_Enzyme_number_p, 1.0)

Versions:

(NonDecaMWE) pkg> st
Status `~/Documents/Work/dev/DecapodesAutoDiff/NonDecaMWE/Project.toml`
  [b0b7db55] ComponentArrays v0.15.13
  [7da242da] Enzyme v0.12.9
  [1dea7af3] OrdinaryDiffEq v6.80.1
  [d236fae5] PreallocationTools v0.4.22
  [1ed8b502] SciMLSensitivity v7.60.1
  [e88e6eb3] Zygote v0.6.70

jClugstor commented 1 month ago

Errors: Zygote:

ERROR: ArgumentError: new: too few arguments (expected 49)
Stacktrace:
  [1] __new__(::Type, ::ODESolution{…}, ::Vararg{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/tools/builtins.jl:9
  [2] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:296 [inlined]
  [3] adjoint(::Zygote.Context{…}, ::typeof(Zygote.__new__), ::Type, ::ODESolution{…}, ::ComponentVector{…}, ::Nothing, ::Vector{…}, ::Vararg{…})
    @ Zygote ./none:0
  [4] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
  [5] ODEIntegrator
    @ ~/.julia/packages/OrdinaryDiffEq/tAI61/src/integrators/type.jl:168 [inlined]
  [6] _pullback(::Zygote.Context{…}, ::Type{…}, ::ODESolution{…}, ::ComponentVector{…}, ::Nothing, ::Vector{…}, ::Float64, ::Float64, ::ODEFunction{…}, ::@NamedTuple{…}, ::ComponentVector{…}, ::ComponentVector{…}, ::Nothing, ::Float64, ::FBDF{…}, ::Float64, ::Bool, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Int64, ::Int64, ::Int64, ::Int64, ::OrdinaryDiffEq.FBDFCache{…}, ::Nothing, ::Int64, ::Bool, ::Bool, ::Bool, ::Bool, ::Int64, ::Int64, ::Float64, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::OrdinaryDiffEq.DEOptions{…}, ::SciMLBase.DEStats, ::OrdinaryDiffEq.DefaultInit, ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
  [7] #__init#806
    @ ~/.julia/packages/OrdinaryDiffEq/tAI61/src/solve.jl:468 [inlined]
  [8] _pullback(::Zygote.Context{…}, ::OrdinaryDiffEq.var"##__init#806", ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Nothing, ::Bool, ::Bool, ::Bool, ::Nothing, ::Nothing, ::Bool, ::Bool, ::Float64, ::Float64, ::Float64, ::Bool, ::Bool, ::Rational{…}, ::Nothing, ::Nothing, ::Rational{…}, ::Int64, ::Rational{…}, ::Rational{…}, ::Nothing, ::Nothing, ::Rational{…}, ::Nothing, ::Bool, ::Int64, ::Int64, ::typeof(DiffEqBase.ODE_DEFAULT_NORM), ::typeof(LinearAlgebra.opnorm), ::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), ::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Int64, ::String, ::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), ::Symbol, ::Nothing, ::Bool, ::Bool, ::Bool, ::Bool, ::OrdinaryDiffEq.DefaultInit, ::@Kwargs{}, ::typeof(SciMLBase.__init), ::ODEProblem{…}, ::FBDF{…}, ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Type{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
  [9] __init
    @ ~/.julia/packages/OrdinaryDiffEq/tAI61/src/solve.jl:11 [inlined]
 [10] _pullback(::Zygote.Context{…}, ::typeof(SciMLBase.__init), ::ODEProblem{…}, ::FBDF{…}, ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Type{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [11] __init (repeats 4 times)
    @ ~/.julia/packages/OrdinaryDiffEq/tAI61/src/solve.jl:11 [inlined]
 [12] _apply
    @ ./boot.jl:838 [inlined]
 [13] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [14] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [15] #__solve#805
    @ ~/.julia/packages/OrdinaryDiffEq/tAI61/src/solve.jl:6 [inlined]
 [16] _pullback(::Zygote.Context{…}, ::OrdinaryDiffEq.var"##__solve#805", ::@Kwargs{}, ::typeof(SciMLBase.__solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [17] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [18] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [19] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [20] __solve
    @ ~/.julia/packages/OrdinaryDiffEq/tAI61/src/solve.jl:1 [inlined]
 [21] _pullback(::Zygote.Context{…}, ::typeof(SciMLBase.__solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [22] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [23] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [24] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [25] #solve_call#44
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:612 [inlined]
 [26] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve_call#44", ::Bool, ::Nothing, ::@Kwargs{}, ::typeof(DiffEqBase.solve_call), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [27] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [28] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [29] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [30] solve_call
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:569 [inlined]
 [31] _pullback(::Zygote.Context{…}, ::typeof(DiffEqBase.solve_call), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [32] #solve_up#53
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1080 [inlined]
 [33] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve_up#53", ::@Kwargs{}, ::typeof(DiffEqBase.solve_up), ::ODEProblem{…}, ::ZygoteVJP, ::ComponentVector{…}, ::@NamedTuple{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [34] _apply
    @ ./boot.jl:838 [inlined]
 [35] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [36] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [37] solve_up
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1066 [inlined]
 [38] _pullback(::Zygote.Context{…}, ::typeof(DiffEqBase.solve_up), ::ODEProblem{…}, ::ZygoteVJP, ::ComponentVector{…}, ::@NamedTuple{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [39] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [40] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [41] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [42] #solve#51
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1003 [inlined]
 [43] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve#51", ::ZygoteVJP, ::Nothing, ::Nothing, ::Val{…}, ::@Kwargs{}, ::typeof(solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [44] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [45] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [46] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [47] solve
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:993 [inlined]
 [48] _pullback(::Zygote.Context{…}, ::typeof(Core.kwcall), ::@NamedTuple{…}, ::typeof(solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [49] fake_loss_Zygote
    @ ~/Documents/Work/dev/DecapodesAutoDiff/NonDecaMWE/oscillatorcalibrate_no_decapode.jl:65 [inlined]
 [50] _pullback(ctx::Zygote.Context{false}, f::typeof(fake_loss_Zygote), args::@NamedTuple{k::Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [51] pullback(f::Function, cx::Zygote.Context{false}, args::@NamedTuple{k::Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:90
 [52] pullback
    @ ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:88 [inlined]
 [53] gradient(f::Function, args::@NamedTuple{k::Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:147
Some type information was truncated. Use `show(err)` to see complete types.

Enzyme:

ERROR: Enzyme execution failed.
Mismatched activity for:   store i64 %11, i64 addrspace(11)* %10, align 8, !dbg !42 const val:   %11 = load i64, i64 addrspace(11)* %9, align 8, !dbg !42
Type tree: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}
 llvalue=  %11 = load i64, i64 addrspace(11)* %9, align 8, !dbg !42
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] _
   @ ~/.julia/packages/SciMLBase/JUp1I/src/problems/ode_problems.jl:118
 [2] ODEProblem
   @ ~/.julia/packages/SciMLBase/JUp1I/src/problems/ode_problems.jl:111
 [3] #remake#653
   @ ~/.julia/packages/SciMLBase/JUp1I/src/remake.jl:133
 [4] remake
   @ ~/.julia/packages/SciMLBase/JUp1I/src/remake.jl:83
 [5] remake
   @ ~/.julia/packages/SciMLBase/JUp1I/src/remake.jl:0

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/F71IJ/src/compiler.jl:1338
  [2] _
    @ ~/.julia/packages/SciMLBase/JUp1I/src/problems/ode_problems.jl:118 [inlined]
  [3] ODEProblem
    @ ~/.julia/packages/SciMLBase/JUp1I/src/problems/ode_problems.jl:111 [inlined]
  [4] #remake#653
    @ ~/.julia/packages/SciMLBase/JUp1I/src/remake.jl:133 [inlined]
  [5] remake
    @ ~/.julia/packages/SciMLBase/JUp1I/src/remake.jl:83 [inlined]
  [6] remake
    @ ~/.julia/packages/SciMLBase/JUp1I/src/remake.jl:0 [inlined]
  [7] augmented_julia_remake_20796_inner_1wrap
    @ ~/.julia/packages/SciMLBase/JUp1I/src/remake.jl:0
  [8] macro expansion
    @ ~/.julia/packages/Enzyme/F71IJ/src/compiler.jl:5916 [inlined]
  [9] enzyme_call
    @ ~/.julia/packages/Enzyme/F71IJ/src/compiler.jl:5566 [inlined]
 [10] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/F71IJ/src/compiler.jl:5454 [inlined]
 [11] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(Core.kwcall), df::Nothing, primal_1::@NamedTuple{…}, shadow_1_1::Base.RefValue{…}, primal_2::typeof(remake), shadow_2_1::Nothing, primal_3::ODEProblem{…}, shadow_3_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/F71IJ/src/rules/jitrules.jl:179
 [12] fake_loss_Zygote
    @ ~/Documents/Work/dev/DecapodesAutoDiff/NonDecaMWE/oscillatorcalibrate_no_decapode.jl:64
 [13] call_composed
    @ ./operators.jl:1044 [inlined]
 [14] ComposedFunction
    @ ./operators.jl:1041 [inlined]
 [15] augmented_julia_ComposedFunction_20348wrap
    @ ./operators.jl:0
 [16] macro expansion
    @ ~/.julia/packages/Enzyme/F71IJ/src/compiler.jl:5916 [inlined]
 [17] enzyme_call
    @ ~/.julia/packages/Enzyme/F71IJ/src/compiler.jl:5566 [inlined]
 [18] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/F71IJ/src/compiler.jl:5454 [inlined]
 [19] autodiff
    @ ~/.julia/packages/Enzyme/F71IJ/src/Enzyme.jl:235 [inlined]
 [20] autodiff
    @ ~/.julia/packages/Enzyme/F71IJ/src/Enzyme.jl:303 [inlined]
 [21] gradient(rm::ReverseMode{false, FFIABI, false}, f::typeof(fake_loss_Zygote), x::@NamedTuple{k::Float64})
    @ Enzyme ~/.julia/packages/Enzyme/F71IJ/src/Enzyme.jl:977
 [22] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/NonDecaMWE/oscillatorcalibrate_no_decapode.jl:83
Some type information was truncated. Use `show(err)` to see complete types.

jClugstor commented 1 month ago

I just realized that in the original version where the gradient worked, the parameters had to be a ComponentArray. However the Decapodes folks use NamedTuples for their parameters. ComponentArrays seem to work to simulate anyway though.

Even with the parameter as a component array, this example still fails.

using ComponentArrays
using PreallocationTools
using OrdinaryDiffEq

not_decapode_f = begin

    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
            end
    end
end

x0 = [1.0]
v0 = [0.0]

u₀ = ComponentArray(X = x0, V = v0)
p = ComponentArray(k = 1.0)
tₑ = 10

fₘ = not_decapode_f()

data_prob = ODEProblem{true, SciMLBase.FullSpecialize}(fₘ, u₀, (0, tₑ),p)
sol = solve(data_prob,Tsit5(), saveat = 0.1)
dat = vcat([u.X for u in sol.u]...)

using SciMLSensitivity, Zygote, Enzyme

function fake_loss_default(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, Tsit5(), sensealg = ZygoteVJP())
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss_default, p)
Enzyme.gradient(Enzyme.Reverse, fake_loss_default, p)
Enzyme.gradient(Enzyme.Reverse, fake_loss_default, p)

function fake_loss_Zygote(p)
  prob = remake(data_prob, p = p)
  soln = solve(prob, FBDF(autodiff = false), sensealg = ZygoteVJP())
  @info("Done")

  # return soln(tₑ)
  sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

function fake_loss_Enzyme(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, FBDF(autodiff = false), sensealg = EnzymeVJP())
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss_Zygote, p)
Zygote.gradient(fake_loss_Enzyme, p)
Enzyme.gradient(Enzyme.Reverse, fake_loss_Zygote, p)
Enzyme.gradient(Enzyme.Reverse, fake_loss_Enzyme, p)

jpfairbanks commented 1 month ago

We could use ComponentArrays for Parameters if necessary. We just need to be able to index by symbols and store heterogeneous collections

ChrisRackauckas commented 1 month ago

Can you isolate this down to just not_decapode_f? Enzyme vs Zygote doesn't matter because it's just hitting this code https://github.com/SciML/SciMLSensitivity.jl/blob/master/src/derivative_wrappers.jl#L702-L706 and so it's about reproducing the error of that, which is just differnetiation of f.

jClugstor commented 1 month ago

You mean something like this?

using ComponentArrays
using PreallocationTools
using Zygote
using Enzyme

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
            end
    end
end

du = ComponentArray(X = [1.0], V = [0.0])
u = ComponentArray(X = [1.0], V = [0.0])
p = ComponentArray(k = 1.0)
f = not_decapode_f()
f(du,u,p,1.0)

Zygote.jacobian(f, du,u,p,1.0)
Enzyme.gradient(Enzyme.Reverse, f, [du,u,p,1.0])

Zygote error:

ERROR: Mutating arrays is not supported -- called copyto!(SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, ...)
This error occurs when you ask Zygote to differentiate operations that change
the elements of arrays in place (e.g. setting values with x .= ...)

Possible fixes:
- avoid mutating operations (preferred)
- or read the documentation and solutions for this error
  https://fluxml.ai/Zygote.jl/latest/limitations

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] _throw_mutation_error(f::Function, args::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/lib/array.jl:70
  [3] (::Zygote.var"#543#544"{SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}})(::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/lib/array.jl:85
  [4] (::Zygote.var"#2633#back#545"{Zygote.var"#543#544"{SubArray{…}}})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
  [5] materialize!
    @ ./broadcast.jl:914 [inlined]
  [6] materialize!
    @ ./broadcast.jl:911 [inlined]
  [7] materialize!
    @ ./broadcast.jl:907 [inlined]
  [8] f
    @ ~/Documents/Work/dev/DecapodesAutoDiff/NonDecaMWE/not_decapode.jl:23 [inlined]
  [9] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [10] (::Zygote.var"#291#292"{Tuple{NTuple{…}}, Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:206
 [11] (::Zygote.var"#2169#back#293"{Zygote.var"#291#292"{Tuple{…}, Zygote.Pullback{…}}})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [12] call_composed
    @ ./operators.jl:1045 [inlined]
 [13] (::Zygote.Pullback{Tuple{typeof(Base.call_composed), Tuple{…}, Tuple{…}, @Kwargs{}}, Any})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [14] call_composed
    @ ./operators.jl:1044 [inlined]
 [15] #_#103
    @ ./operators.jl:1041 [inlined]
 [16] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [17] #291
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:206 [inlined]
 [18] #2169#back
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72 [inlined]
 [19] ComposedFunction
    @ ./operators.jl:1041 [inlined]
 [20] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [21] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:91
 [22] withjacobian(::Function, ::ComponentVector{…}, ::ComponentVector{…}, ::Vararg{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/lib/grad.jl:150
 [23] jacobian(::Function, ::ComponentVector{…}, ::ComponentVector{…}, ::Vararg{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/lib/grad.jl:128
 [24] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/NonDecaMWE/not_decapode.jl:34
Some type information was truncated. Use `show(err)` to see complete types.

Enzyme error:

ERROR: UndefRefError: access to undefined reference
Stacktrace:
  [1] getproperty
    @ ./Base.jl:37 [inlined]
  [2] getindex
    @ ./refvalue.jl:59 [inlined]
  [3] unsafe_function_from_type
    @ ~/.julia/packages/GPUCompiler/nWT2N/src/jlgen.jl:32 [inlined]
  [4] MethodError
    @ ~/.julia/packages/GPUCompiler/nWT2N/src/jlgen.jl:36 [inlined]
  [5] macro expansion
    @ ~/.julia/packages/Enzyme/F71IJ/src/utils.jl:0 [inlined]
  [6] codegen_world_age(ft::Type{var"#f#7"{FixedSizeDiffCache{…}, FixedSizeDiffCache{…}}}, tt::Type{Tuple{Vector{…}}})
    @ Enzyme ~/.julia/packages/Enzyme/F71IJ/src/utils.jl:168
  [7] autodiff
    @ ~/.julia/packages/Enzyme/F71IJ/src/Enzyme.jl:224 [inlined]
  [8] autodiff
    @ ~/.julia/packages/Enzyme/F71IJ/src/Enzyme.jl:303 [inlined]
  [9] gradient(rm::ReverseMode{…}, f::var"#f#7"{…}, x::Vector{…})
    @ Enzyme ~/.julia/packages/Enzyme/F71IJ/src/Enzyme.jl:981
 [10] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/NonDecaMWE/not_decapode.jl:35
Some type information was truncated. Use `show(err)` to see complete types.

ChrisRackauckas commented 1 month ago

using ComponentArrays
using FiniteDiff
using Enzyme

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = Vector{Float64}(undef, 1)
            __V̇ = Vector{Float64}(undef, 1)
        end
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = var"__•1"
                V̇ = __V̇
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
end

du = ComponentArray(X = [0.0], V = [0.0])
u = ComponentArray(X = [1.0], V = [3.0])
p = ComponentArray(k = 1.0)
f = not_decapode_f()
f(du,u,p,1.0)
df = Enzyme.make_zero(f)
d_du = Enzyme.make_zero(du)
d_u = Enzyme.make_zero(u)
dp = Enzyme.make_zero(p)

d_du .= 0; d_u .= 0; dp[1] = 1.0
Enzyme.autodiff(Enzyme.Forward, Duplicated(f, df), Enzyme.Duplicated(du, d_du),
                Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp),
                Enzyme.Const(1.0))

d_du # ComponentVector{Float64}(X = [0.0], V = [-1.0])
dp # ComponentVector{Float64}(k = 1.0)
du # ComponentVector{Float64}(X = [3.0], V = [-1.0])

d_du .= 0; d_u .= 0; dp[1] = 1.0
Enzyme.autodiff(Enzyme.Reverse, Duplicated(f, df), Enzyme.Duplicated(du, d_du),
                Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp),
                Enzyme.Const(1.0))

d_du # ComponentVector{Float64}(X = [0.0], V = [0.0])
dp # ComponentVector{Float64}(k = 3.0)
du # ComponentVector{Float64}(X = [3.0], V = [-1.0])

# Confirm with finite difference

ppet = ComponentArray(k = 1 + 1e-7)
du2 = copy(du)
f(du,u,p,1.0)
f(du2,u,ppet,1.0)
(du2 - du) ./ 1e-7 # ComponentVector{Float64}(X = [0.0], V = [-1.0000000005838672])

Interesting issue, but target this form.

ChrisRackauckas commented 1 month ago

Well the PreallocationTools is still used for the ForwardDiff in the ODE solver. But the reverse mode shouldn't care.

wsmoses commented 4 weeks ago

using Enzyme

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = Vector{Float64}(undef, 1)
            __V̇ = Vector{Float64}(undef, 1)
        end
        f(du, u, p, t) = begin
                begin
                    X = u[1]
                    k = p[1]
                    V = u[2]
                end
                var"•1" = var"__•1"
                V̇ = __V̇
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                du[1] = V
                du[2] = V̇[1]
                nothing
            end
    end
end

du = [0.0,0.0]
u = [1.0,3.0]
p = [1.0]
f = not_decapode_f()
f(du,u,p,1.0)
df = Enzyme.make_zero(f)
d_du = Enzyme.make_zero(du)
d_u = Enzyme.make_zero(u)
dp = Enzyme.make_zero(p)

df = Enzyme.make_zero(f)
d_du .= 0; d_u .= 0; dp[1] = 1.0
Enzyme.autodiff(Enzyme.Reverse, Duplicated(f, df), Enzyme.Duplicated(du, d_du),
                Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp),
                Enzyme.Const(1.0))

@show d_du # [0.0,0.0]
@show dp # [3.0]
@show du # [3.0,-1.0]

df = Enzyme.make_zero(f)
# compute the gradient wrt d_u[1]
d_du = [0.0, 1.0]; d_u .= 0; dp[1] = 0.0
Enzyme.autodiff(Enzyme.Reverse, Duplicated(f, df), Enzyme.Duplicated(du, d_du),
                Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp),
                Enzyme.Const(1.0))

# derivative of d_u[1] / dp, which is what finite differences computes below [in the second term].
@show dp # dp = [-1.0]

# Confirm with finite difference

ppet = [1 + 1e-7]
du2 = copy(du)
f(du,u,p,1.0)
f(du2,u,ppet,1.0)
@show (du2 - du) ./ 1e-7 # [0.0,-1.0000000005838672]

wsmoses commented 4 weeks ago

whoops commented in the wrong issue. but in any case @ChrisRackauckas your issue here was calling Enzyme incorrectly.

In reverse mode you need to set the shadow of the return to 1, not the shadow of the input.

The second case above shows the correct result when called properly

ChrisRackauckas commented 4 weeks ago

Yup, dope.

So then this is exactly what's needed in the context of the real model:

using Enzyme, PreallocationTools, ComponentArrays

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
end

du = ComponentArray(X = [0.0], V = [0.0])
u = ComponentArray(X = [1.0], V = [3.0])
p = ComponentArray(k = 1.0)
f = not_decapode_f()
f(du,u,p,1.0)
df = Enzyme.make_zero(f)
d_du = Enzyme.make_zero(du)
d_u = Enzyme.make_zero(u)
dp = Enzyme.make_zero(p)

d_du .= 0; d_u .= 0; dp[1] = 1.0
Enzyme.autodiff(Enzyme.Forward, Duplicated(f, df), Enzyme.Duplicated(du, d_du),
                Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp),
                Enzyme.Const(1.0))

d_du # ComponentVector{Float64}(X = [0.0], V = [-1.0])
dp # ComponentVector{Float64}(k = 1.0)
du # ComponentVector{Float64}(X = [3.0], V = [-1.0])

d_du .= 0; d_u .= 0; dp[1] = 0.0
d_du = ComponentArray(X = [0.0], V = [1.0])
Enzyme.autodiff(Enzyme.Reverse, Duplicated(f, df), Enzyme.Duplicated(du, d_du),
                Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp),
                Enzyme.Const(1.0))

d_du # ComponentVector{Float64}(X = [0.0], V = [0.0])
dp # ComponentVector{Float64}(k = 3.0)
du # ComponentVector{Float64}(X = [3.0], V = [-1.0])

# Confirm with finite difference

ppet = ComponentArray(k = 1 + 1e-7)
du2 = copy(du)
f(du,u,p,1.0)
f(du2,u,ppet,1.0)
(du2 - du) ./ 1e-7 # ComponentVector{Float64}(X = [0.0], V = [-1.0000000005838672])

which works on my machine. @jClugstor can you make our code target that?

wsmoses commented 4 weeks ago

To be clear you don't need to reallocate a new array for d_du you just need to set the values to 0 except 1 where you want the derivatives

On Fri, Jun 7, 2024, 2:03 AM Christopher Rackauckas < @.***> wrote:

Yup, dope.

So then this is exactly what's needed in the context of the real model:

using Enzyme, PreallocationTools, ComponentArrays

not_decapode_f = begin function simulate() begin var"•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1)) V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1)) end f(du, u, p, t) = begin begin X = u.X k = p.k V = u.V end var"•1" = PreallocationTools.get_tmp(var"__•1", u) V̇ = PreallocationTools.get_tmp(__V̇, u) var"•1" .= (.-)(k) V̇ .= var"•1" .* X getproperty(du, :X) .= V getproperty(du, :V) .= V̇ nothing end endend

du = ComponentArray(X = [0.0], V = [0.0]) u = ComponentArray(X = [1.0], V = [3.0]) p = ComponentArray(k = 1.0) f = not_decapode_f()f(du,u,p,1.0) df = Enzyme.make_zero(f) d_du = Enzyme.make_zero(du) d_u = Enzyme.make_zero(u) dp = Enzyme.make_zero(p)

d_du .= 0; d_u .= 0; dp[1] = 1.0 Enzyme.autodiff(Enzyme.Forward, Duplicated(f, df), Enzyme.Duplicated(du, d_du), Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp), Enzyme.Const(1.0))

d_du # ComponentVector{Float64}(X = [0.0], V = [-1.0]) dp # ComponentVector{Float64}(k = 1.0) du # ComponentVector{Float64}(X = [3.0], V = [-1.0])

d_du .= 0; d_u .= 0; dp[1] = 0.0 d_du = ComponentArray(X = [0.0], V = [1.0]) Enzyme.autodiff(Enzyme.Reverse, Duplicated(f, df), Enzyme.Duplicated(du, d_du), Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp), Enzyme.Const(1.0))

d_du # ComponentVector{Float64}(X = [0.0], V = [0.0]) dp # ComponentVector{Float64}(k = 3.0) du # ComponentVector{Float64}(X = [3.0], V = [-1.0])

Confirm with finite difference

ppet = ComponentArray(k = 1 + 1e-7) du2 = copy(du)f(du,u,p,1.0)f(du2,u,ppet,1.0) (du2 - du) ./ 1e-7 # ComponentVector{Float64}(X = [0.0], V = [-1.0000000005838672])

which works on my machine. @jClugstor https://github.com/jClugstor can you make our code target that?

— Reply to this email directly, view it on GitHub https://github.com/DARPA-ASKEM/sciml-service/issues/177#issuecomment-2153604811, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXE34VNCD2SW66TYAY3ZGD2FJAVCNFSM6AAAAABIR6QLPWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCNJTGYYDIOBRGE . You are receiving this because you commented.Message ID: @.***>

ChrisRackauckas commented 4 weeks ago

Yeah this is just easiest to write.

Note that SciMLSensitivity probably needs the make_zero! on the df, I can add that when it's added to Enzyme. Technically no one should be caching anything in there that isn't written to on a call, but we can improve the safety anyways.

jClugstor commented 3 weeks ago

The only thing preventing the Decapode generated code from being autodiffed with Enzyme now seems to be that the generated function needs to return nothing . @jpfairbanks is there an option for that in Decapodes?

ChrisRackauckas commented 3 weeks ago

Just add it in there.

https://github.com/AlgebraicJulia/Decapodes.jl/blob/main/src/simulation.jl#L565-L576

It's one line right there.

jClugstor commented 3 weeks ago

Oh cool thanks. In the meantime I just put it there manually.

cd(@__DIR__)
using Pkg
Pkg.activate(".")
Pkg.instantiate()

using Catlab
using Catlab.Graphics
using CombinatorialSpaces
using ComponentArrays
using Catlab.Graphics
using CombinatorialSpaces.ExteriorCalculus
using DiagrammaticEquations
using DiagrammaticEquations.Deca
using Decapodes
using GeometryBasics:Point1,Point2,Point3
using OrdinaryDiffEq
using Enzyme
using SciMLSensitivity

Point1D = Point1{Float64}
Point2D = Point2{Float64};
Point3D = Point3{Float64};

decapode_f = begin
    #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:570 =#
    function simulate(mesh, operators, hodge = GeometricHodge())
        #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:570 =#
        #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:571 =#
        begin
            #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:174 =#
        end
        #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:572 =#
        begin
            #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:468 =#
        end
        #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:573 =#
        begin
            #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:227 =#
            var"__•1" = Decapodes.FixedSizeDiffCache(Vector{Float64}(undef, nparts(mesh, :V)))
            __V̇ = Decapodes.FixedSizeDiffCache(Vector{Float64}(undef, nparts(mesh, :V)))
        end
        #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:574 =#
        f(du, u, p, t) = begin
                #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:574 =#
                #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:575 =#
                begin
                    #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:252 =#
                    X = u.X
                    k = p.k
                    V = u.V
                end
                #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:576 =#
                var"•1" = Decapodes.get_tmp(var"__•1", u)
                V̇ = Decapodes.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                #= /home/jadonclugston/.julia/packages/Decapodes/qxJAY/src/simulation.jl:577 =#
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
end

x0 = [1.0]
v0 = [0.0]

u₀ = ComponentArray(X = x0, V = v0)
p = ComponentArray(k = 1.0)
tₑ = 10

s_prime = EmbeddedDeltaSet1D{Bool, Point1D}()
add_vertices!(s_prime,1)

s = EmbeddedDeltaDualComplex1D{Bool, Float64, Point1D}(s_prime)

function generate(operators,mesh)
end

f = decapode_f(s, generate)

#Enzyme autodiff f test ============================================================================================================================
du = ComponentArray(X = [0.0], V = [0.0])
u = ComponentArray(X = [1.0], V = [3.0])
p = ComponentArray(k = 1.0)
f(du,u,p,1.0)
df = Enzyme.make_zero(f)
d_du = Enzyme.make_zero(du)
d_u = Enzyme.make_zero(u)
dp = Enzyme.make_zero(p)

d_du .= 0; d_u .= 0; dp[1] = 1.0
Enzyme.autodiff(Enzyme.Forward, Duplicated(f, df), Enzyme.Duplicated(du, d_du),
                Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp),
                Enzyme.Const(1.0))

d_du # ComponentVector{Float64}(X = [0.0], V = [-1.0])
dp # ComponentVector{Float64}(k = 1.0)
du # ComponentVector{Float64}(X = [3.0], V = [-1.0])

d_du .= 0; d_u .= 0; dp[1] = 0.0
d_du = ComponentArray(X = [0.0], V = [1.0])
Enzyme.autodiff(Enzyme.Reverse, Duplicated(f, df), Enzyme.Duplicated(du, d_du),
                Enzyme.Duplicated(u,d_u), Enzyme.Duplicated(p,dp),
                Enzyme.Const(1.0))

d_du # ComponentVector{Float64}(X = [0.0], V = [0.0])
dp # ComponentVector{Float64}(k = 3.0)
du # ComponentVector{Float64}(X = [3.0], V = [-1.0])
# ===================================================================================================================================================

data_prob = ODEProblem{true, SciMLBase.FullSpecialize}(f, u₀, (0, tₑ),p)
sol = solve(data_prob,Tsit5(), saveat = 0.1)
dat = vcat([u.X for u in sol.u]...)

function fake_loss(p)
  prob = remake(data_prob, p = p)
  soln = solve(prob, FBDF(autodiff = false), sensealg = EnzymeVJP())
  @info("Done")

  # return soln(tₑ)
  sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end
Enzyme.gradient(Enzyme.Reverse, fake_loss, p)

ERROR: Enzyme execution failed.
Mismatched activity for:   store i64 %11, i64 addrspace(11)* %10, align 8, !dbg !42 const val:   %11 = load i64, i64 addrspace(11)* %9, align 8, !dbg !42
Type tree: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}
 llvalue=  %11 = load i64, i64 addrspace(11)* %9, align 8, !dbg !42
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] _
   @ ~/.julia/packages/SciMLBase/roUH9/src/problems/ode_problems.jl:118
 [2] ODEProblem
   @ ~/.julia/packages/SciMLBase/roUH9/src/problems/ode_problems.jl:111
 [3] #remake#653
   @ ~/.julia/packages/SciMLBase/roUH9/src/remake.jl:133
 [4] remake
   @ ~/.julia/packages/SciMLBase/roUH9/src/remake.jl:83
 [5] remake
   @ ~/.julia/packages/SciMLBase/roUH9/src/remake.jl:0

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:1539
  [2] _
    @ ~/.julia/packages/SciMLBase/roUH9/src/problems/ode_problems.jl:118 [inlined]
  [3] ODEProblem
    @ ~/.julia/packages/SciMLBase/roUH9/src/problems/ode_problems.jl:111 [inlined]
  [4] #remake#653
    @ ~/.julia/packages/SciMLBase/roUH9/src/remake.jl:133 [inlined]
  [5] remake
    @ ~/.julia/packages/SciMLBase/roUH9/src/remake.jl:83 [inlined]
  [6] remake
    @ ~/.julia/packages/SciMLBase/roUH9/src/remake.jl:0 [inlined]
  [7] augmented_julia_remake_16021_inner_1wrap
    @ ~/.julia/packages/SciMLBase/roUH9/src/remake.jl:0
  [8] macro expansion
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:6117 [inlined]
  [9] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Val{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Const{…}, ::Const{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5767
 [10] (::Enzyme.Compiler.AugmentedForwardThunk{…})(::Const{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5655
 [11] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(Core.kwcall), df::Nothing, primal_1::@NamedTuple{…}, shadow_1_1::@NamedTuple{…}, primal_2::typeof(remake), shadow_2_1::Nothing, primal_3::ODEProblem{…}, shadow_3_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/rules/jitrules.jl:213
 [12] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeSensitivityMWE.jl:120 [inlined]
 [13] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeSensitivityMWE.jl:0 [inlined]
 [14] augmented_julia_fake_loss_16049_inner_1wrap
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeSensitivityMWE.jl:0
 [15] macro expansion
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:6117 [inlined]
 [16] enzyme_call
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5767 [inlined]
 [17] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5655 [inlined]
 [18] autodiff
    @ ~/.julia/packages/Enzyme/wg4P7/src/Enzyme.jl:242 [inlined]
 [19] autodiff
    @ ~/.julia/packages/Enzyme/wg4P7/src/Enzyme.jl:310 [inlined]
 [20] gradient(rm::ReverseMode{false, FFIABI, false}, f::typeof(fake_loss), x::ComponentVector{Float64, Vector{…}, Tuple{…}})
    @ Enzyme ~/.julia/packages/Enzyme/wg4P7/src/Enzyme.jl:988
 [21] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeSensitivityMWE.jl:127
Some type information was truncated. Use `show(err)` to see complete types.

Looks like autodiff through f is working, but there's something in remake. This was using the branch in https://github.com/SciML/SciMLSensitivity.jl/pull/1057

ChrisRackauckas commented 3 weeks ago

You only need Enzyme for the inside. Just do Zygote on the outside.

jClugstor commented 3 weeks ago

Do you mean sensealg = EnzymeVJP(), but do Zygote.gradient(fake_loss,p) ?

function fake_loss(p)
  prob = remake(data_prob, p = p)
  soln = solve(prob, FBDF(autodiff = false), sensealg = EnzymeVJP())
  @info("Done")

  # return soln(tₑ)
  sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss, p)

ERROR: ArgumentError: new: too few arguments (expected 49)
Stacktrace:
  [1] __new__(::Type, ::ODESolution{…}, ::Vararg{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/tools/builtins.jl:9
  [2] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:296 [inlined]
  [3] adjoint(::Zygote.Context{…}, ::typeof(Zygote.__new__), ::Type, ::ODESolution{…}, ::ComponentVector{…}, ::Nothing, ::Vector{…}, ::Vararg{…})
    @ Zygote ./none:0
  [4] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
  [5] ODEIntegrator
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/integrators/type.jl:168 [inlined]
  [6] _pullback(::Zygote.Context{…}, ::Type{…}, ::ODESolution{…}, ::ComponentVector{…}, ::Nothing, ::Vector{…}, ::Float64, ::Float64, ::ODEFunction{…}, ::ComponentVector{…}, ::ComponentVector{…}, ::ComponentVector{…}, ::Nothing, ::Float64, ::FBDF{…}, ::Float64, ::Bool, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Int64, ::Int64, ::Int64, ::Int64, ::OrdinaryDiffEq.FBDFCache{…}, ::Nothing, ::Int64, ::Bool, ::Bool, ::Bool, ::Bool, ::Int64, ::Int64, ::Float64, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::OrdinaryDiffEq.DEOptions{…}, ::SciMLBase.DEStats, ::OrdinaryDiffEq.DefaultInit, ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
  [7] #__init#671
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:474 [inlined]
  [8] _pullback(::Zygote.Context{…}, ::OrdinaryDiffEq.var"##__init#671", ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Nothing, ::Bool, ::Bool, ::Bool, ::Nothing, ::Nothing, ::Bool, ::Bool, ::Float64, ::Float64, ::Float64, ::Bool, ::Bool, ::Rational{…}, ::Nothing, ::Nothing, ::Rational{…}, ::Int64, ::Rational{…}, ::Rational{…}, ::Nothing, ::Nothing, ::Rational{…}, ::Nothing, ::Bool, ::Int64, ::Int64, ::typeof(DiffEqBase.ODE_DEFAULT_NORM), ::typeof(LinearAlgebra.opnorm), ::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), ::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Int64, ::String, ::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), ::Symbol, ::Nothing, ::Bool, ::Bool, ::Bool, ::Bool, ::OrdinaryDiffEq.DefaultInit, ::@Kwargs{}, ::typeof(SciMLBase.__init), ::ODEProblem{…}, ::FBDF{…}, ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Type{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
  [9] __init
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:11 [inlined]
 [10] _pullback(::Zygote.Context{…}, ::typeof(SciMLBase.__init), ::ODEProblem{…}, ::FBDF{…}, ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Type{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [11] __init (repeats 4 times)
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:11 [inlined]
 [12] _apply
    @ ./boot.jl:838 [inlined]
 [13] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [14] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [15] #__solve#670
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:6 [inlined]
 [16] _pullback(::Zygote.Context{…}, ::OrdinaryDiffEq.var"##__solve#670", ::@Kwargs{}, ::typeof(SciMLBase.__solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [17] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [18] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [19] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [20] __solve
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:1 [inlined]
 [21] _pullback(::Zygote.Context{…}, ::typeof(SciMLBase.__solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [22] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [23] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [24] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [25] #solve_call#44
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:612 [inlined]
 [26] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve_call#44", ::Bool, ::Nothing, ::@Kwargs{}, ::typeof(DiffEqBase.solve_call), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [27] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [28] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [29] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [30] solve_call
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:569 [inlined]
 [31] _pullback(::Zygote.Context{…}, ::typeof(DiffEqBase.solve_call), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [32] #solve_up#53
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1080 [inlined]
 [33] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve_up#53", ::@Kwargs{}, ::typeof(DiffEqBase.solve_up), ::ODEProblem{…}, ::EnzymeVJP, ::ComponentVector{…}, ::ComponentVector{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [34] _apply
    @ ./boot.jl:838 [inlined]
 [35] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [36] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [37] solve_up
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1066 [inlined]
 [38] _pullback(::Zygote.Context{…}, ::typeof(DiffEqBase.solve_up), ::ODEProblem{…}, ::EnzymeVJP, ::ComponentVector{…}, ::ComponentVector{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [39] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [40] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [41] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [42] #solve#51
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1003 [inlined]
 [43] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve#51", ::EnzymeVJP, ::Nothing, ::Nothing, ::Val{…}, ::@Kwargs{}, ::typeof(solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [44] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [45] adjoint
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:203 [inlined]
 [46] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [47] solve
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:993 [inlined]
 [48] _pullback(::Zygote.Context{…}, ::typeof(Core.kwcall), ::@NamedTuple{…}, ::typeof(solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [49] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeSensitivityMWE.jl:121 [inlined]
 [50] _pullback(ctx::Zygote.Context{false}, f::typeof(fake_loss), args::ComponentVector{Float64, Vector{…}, Tuple{…}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [51] pullback(f::Function, cx::Zygote.Context{false}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:90
 [52] pullback
    @ ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:88 [inlined]
 [53] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:147
Some type information was truncated. Use `show(err)` to see complete types.

ChrisRackauckas commented 3 weeks ago

  [1] __new__(::Type, ::ODESolution{…}, ::Vararg{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/tools/builtins.jl:9

Make this print out the args.

jClugstor commented 3 weeks ago

args length: 47
args: (ODESolution{Float64, 2, Vector{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}}}, ODEProblem{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, Tuple{Float64, Float64}, true, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}}, ODEFunction{true, SciMLBase.FullSpecialize, var"#f#7"{PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}, Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}, OrdinaryDiffEq.InterpolationData{ODEFunction{true, SciMLBase.FullSpecialize, var"#f#7"{PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}}, Vector{Float64}, Vector{Vector{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}}}, Nothing, OrdinaryDiffEq.Tsit5Cache{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}, Nothing}, SciMLBase.DEStats, Nothing, Nothing, Nothing}(ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}[], nothing, nothing, Float64[], Vector{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}}[], ODEProblem{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, Tuple{Float64, Float64}, true, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}}, ODEFunction{true, SciMLBase.FullSpecialize, var"#f#7"{PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}(ODEFunction{true, SciMLBase.FullSpecialize, var"#f#7"{PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(f, LinearAlgebra.UniformScaling{Bool}(true), nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), (X = [1.0], V = [0.0]), (0.0, 10.0), (k = 1.0), Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}(), SciMLBase.StandardODEProblem()), Tsit5(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),), OrdinaryDiffEq.InterpolationData{ODEFunction{true, SciMLBase.FullSpecialize, var"#f#7"{PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}}, Vector{Float64}, Vector{Vector{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}}}, Nothing, OrdinaryDiffEq.Tsit5Cache{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}, Nothing}(ODEFunction{true, SciMLBase.FullSpecialize, var"#f#7"{PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(f, LinearAlgebra.UniformScaling{Bool}(true), nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}[], Float64[], Vector{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}}[], nothing, true, OrdinaryDiffEq.Tsit5Cache{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}((X = [1.0], V = [0.0]), (X = [1.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), OrdinaryDiffEq.trivial_limiter!, OrdinaryDiffEq.trivial_limiter!, static(false)), nothing, false), true, 0, SciMLBase.DEStats(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0), nothing, SciMLBase.ReturnCode.Default, nothing, nothing), (X = [1.0], V = [0.0]), nothing, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}[], 0.0, 0.0, ODEFunction{true, SciMLBase.FullSpecialize, var"#f#7"{PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, PreallocationTools.FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(f, LinearAlgebra.UniformScaling{Bool}(true), nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), (k = 1.0), (X = [1.0], V = [0.0]), (X = [1.0], V = [0.0]), nothing, 0.0, Tsit5(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),), 0.0, true, 0.0, 1.0, 1.0, 1.0, 0.0001, 1.0, 1.0, 1.0, 0, 0, 0, 0, OrdinaryDiffEq.Tsit5Cache{ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, typeof(OrdinaryDiffEq.trivial_limiter!), typeof(OrdinaryDiffEq.trivial_limiter!), Static.False}((X = [1.0], V = [0.0]), (X = [1.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), (X = [0.0], V = [0.0]), OrdinaryDiffEq.trivial_limiter!, OrdinaryDiffEq.trivial_limiter!, static(false)), nothing, 0, false, false, false, true, 0, 1, 0.0, false, false, false, false, true, false, OrdinaryDiffEq.DEOptions{Float64, Float64, Float64, Float64, PIController{Rational{Int64}}, typeof(DiffEqBase.ODE_DEFAULT_NORM), typeof(LinearAlgebra.opnorm), Nothing, CallbackSet{Tuple{}, Tuple{}}, typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}, Nothing, Nothing, Int64, Tuple{}, Tuple{}, Tuple{}}(1000000, true, true, 1.0e-6, 0.001, 0.9, 10.0, 0.2, 1.0, 1.0, 0.0001, 2.0, 10.0, 0.0, PIController{Rational{Int64}}(7//50, 2//25), DiffEqBase.ODE_DEFAULT_NORM, LinearAlgebra.opnorm, nothing, DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}(DataStructures.FasterForward(), [10.0]), DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}(DataStructures.FasterForward(), Float64[]), DataStructures.BinaryHeap{Float64, DataStructures.FasterForward}(DataStructures.FasterForward(), Float64[]), (), (), (), nothing, false, 1000, "ODE", DiffEqBase.ODE_DEFAULT_PROG_MESSAGE, :OrdinaryDiffEq, true, false, true, true, true, true, nothing, CallbackSet{Tuple{}, Tuple{}}((), ()), DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN, DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK, true, true, false, false, false), SciMLBase.DEStats(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0), OrdinaryDiffEq.DefaultInit(), nothing)

Looks like there are only 47 args where 49 is expected

jClugstor commented 3 weeks ago

This is a pretty small example that demonstrates the Zygote error

using Enzyme, PreallocationTools, ComponentArrays

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
end

du = ComponentArray(X = [0.0], V = [0.0])
u = ComponentArray(X = [1.0], V = [3.0])
p = ComponentArray(k = 1.0)
f = not_decapode_f()

data_prob = ODEProblem{true, SciMLBase.FullSpecialize}(f, u₀, (0, tₑ),p)
sol = solve(data_prob,Tsit5(), saveat = 0.1)
dat = vcat([u.X for u in sol.u]...)

function fake_loss(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, FBDF(autodiff = false), sensealg = EnzymeVJP())
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss,p)

ERROR: ArgumentError: new: too few arguments (expected 49)
Stacktrace:
  [1] __new__(::Type, ::ODESolution{…}, ::Vararg{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/tools/builtins.jl:14
  [2] adjoint
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:296 [inlined]
  [3] adjoint(::Zygote.Context{…}, ::typeof(Zygote.__new__), ::Type, ::ODESolution{…}, ::ComponentVector{…}, ::Nothing, ::Vector{…}, ::Vararg{…})
    @ Zygote ./none:0
  [4] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
  [5] ODEIntegrator
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/integrators/type.jl:168 [inlined]
  [6] _pullback(::Zygote.Context{…}, ::Type{…}, ::ODESolution{…}, ::ComponentVector{…}, ::Nothing, ::Vector{…}, ::Float64, ::Float64, ::ODEFunction{…}, ::ComponentVector{…}, ::ComponentVector{…}, ::ComponentVector{…}, ::Nothing, ::Float64, ::FBDF{…}, ::Float64, ::Bool, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Float64, ::Int64, ::Int64, ::Int64, ::Int64, ::OrdinaryDiffEq.FBDFCache{…}, ::Nothing, ::Int64, ::Bool, ::Bool, ::Bool, ::Bool, ::Int64, ::Int64, ::Float64, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::OrdinaryDiffEq.DEOptions{…}, ::SciMLBase.DEStats, ::OrdinaryDiffEq.DefaultInit, ::Nothing)
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
  [7] #__init#671
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:474 [inlined]
  [8] _pullback(::Zygote.Context{…}, ::OrdinaryDiffEq.var"##__init#671", ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Nothing, ::Bool, ::Bool, ::Bool, ::Nothing, ::Nothing, ::Bool, ::Bool, ::Float64, ::Float64, ::Float64, ::Bool, ::Bool, ::Rational{…}, ::Nothing, ::Nothing, ::Rational{…}, ::Int64, ::Rational{…}, ::Rational{…}, ::Nothing, ::Nothing, ::Rational{…}, ::Nothing, ::Bool, ::Int64, ::Int64, ::typeof(DiffEqBase.ODE_DEFAULT_NORM), ::typeof(LinearAlgebra.opnorm), ::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), ::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Bool, ::Int64, ::String, ::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), ::Symbol, ::Nothing, ::Bool, ::Bool, ::Bool, ::Bool, ::OrdinaryDiffEq.DefaultInit, ::@Kwargs{}, ::typeof(SciMLBase.__init), ::ODEProblem{…}, ::FBDF{…}, ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Type{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
  [9] __init
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:11 [inlined]
 [10] _pullback(::Zygote.Context{…}, ::typeof(SciMLBase.__init), ::ODEProblem{…}, ::FBDF{…}, ::Tuple{}, ::Tuple{}, ::Tuple{}, ::Type{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [11] __init (repeats 4 times)
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:11 [inlined]
 [12] _apply
    @ ./boot.jl:838 [inlined]
 [13] adjoint
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:203 [inlined]
 [14] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [15] #__solve#670
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:6 [inlined]
 [16] _pullback(::Zygote.Context{…}, ::OrdinaryDiffEq.var"##__solve#670", ::@Kwargs{}, ::typeof(SciMLBase.__solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [17] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [18] adjoint
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:203 [inlined]
 [19] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [20] __solve
    @ ~/.julia/packages/OrdinaryDiffEq/PsydA/src/solve.jl:1 [inlined]
 [21] _pullback(::Zygote.Context{…}, ::typeof(SciMLBase.__solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [22] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [23] adjoint
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:203 [inlined]
 [24] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [25] #solve_call#44
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:612 [inlined]
 [26] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve_call#44", ::Bool, ::Nothing, ::@Kwargs{}, ::typeof(DiffEqBase.solve_call), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [27] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [28] adjoint
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:203 [inlined]
 [29] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [30] solve_call
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:569 [inlined]
 [31] _pullback(::Zygote.Context{…}, ::typeof(DiffEqBase.solve_call), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [32] #solve_up#53
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1080 [inlined]
 [33] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve_up#53", ::@Kwargs{}, ::typeof(DiffEqBase.solve_up), ::ODEProblem{…}, ::EnzymeVJP, ::ComponentVector{…}, ::ComponentVector{…}, ::FBDF{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [34] _apply
    @ ./boot.jl:838 [inlined]
 [35] adjoint
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:203 [inlined]
 [36] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [37] solve_up
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1066 [inlined]
 [38] _pullback(::Zygote.Context{…}, ::typeof(DiffEqBase.solve_up), ::ODEProblem{…}, ::EnzymeVJP, ::ComponentVector{…}, ::ComponentVector{…}, ::FBDF{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [39] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [40] adjoint
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:203 [inlined]
 [41] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [42] #solve#51
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1003 [inlined]
 [43] _pullback(::Zygote.Context{…}, ::DiffEqBase.var"##solve#51", ::EnzymeVJP, ::Nothing, ::Nothing, ::Val{…}, ::@Kwargs{}, ::typeof(solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [44] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:838
 [45] adjoint
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:203 [inlined]
 [46] _pullback
    @ ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:67 [inlined]
 [47] solve
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:993 [inlined]
 [48] _pullback(::Zygote.Context{…}, ::typeof(Core.kwcall), ::@NamedTuple{…}, ::typeof(solve), ::ODEProblem{…}, ::FBDF{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [49] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:37 [inlined]
 [50] _pullback(ctx::Zygote.Context{false}, f::typeof(fake_loss), args::ComponentVector{Float64, Vector{…}, Tuple{…}})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [51] pullback(f::Function, cx::Zygote.Context{false}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:90
 [52] pullback
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:88 [inlined]
 [53] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:147
Some type information was truncated. Use `show(err)` to see complete types.

jClugstor commented 3 weeks ago

using Enzyme, PreallocationTools, ComponentArrays, Zygote, OrdinaryDiffEq, SciMLSensitivity

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
end

u = ComponentArray(X = [1.0], V = [0.0])
p = ComponentArray(k = 1.0)
f = not_decapode_f()

data_prob = ODEProblem{true, SciMLBase.FullSpecialize}(f, u, (0, 20),p)
sol = solve(data_prob,Tsit5(), saveat = 0.1)
dat = vcat([u.X for u in sol.u]...)

function fake_loss(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, Tsit5(), sensealg = EnzymeVJP())
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss,p)

ChrisRackauckas commented 3 weeks ago

solve(prob, Tsit5(), sensealg = EnzymeVJP()) that's not proper syntax.

function fake_loss(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, Tsit5(), sensealg = QuadratureAdjoint(autojacvec=EnzymeVJP()))
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

jClugstor commented 3 weeks ago

So I read the SciMLSensitivity docs wrong... oops. It does change what happens so I thought it was valid.

Anyway, changing to sensealg = QuadratureAdjoint(autojacvec=EnzymeVJP()), it gives an error in the generated function again at the line V̇ .= var"•1" .* X

It's the exact same function that can be differentiated through with Enzyme on it's own.

ERROR: Enzyme execution failed.
Mismatched activity for:   %unbox152.fca.1.0.1.insert.pn.extract.0 = phi {} addrspace(10)* [ %unbox152.fca.0.load, %L192 ], [ %getfield, %L175 ] const val:   %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !31, !tbaa !28, !alias.scope !43, !noalias !46, !nonnull !23, !dereferenceable !51, !align !52
 value=Unknown object of type Vector{Float64}
 llvalue=  %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !31, !tbaa !28, !alias.scope !43, !noalias !46, !nonnull !23, !dereferenceable !51, !align !52
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] <
   @ ./int.jl:83
 [2] macro expansion
   @ ./simdloop.jl:72
 [3] copyto!
   @ ./broadcast.jl:1003
 [4] copyto!
   @ ./broadcast.jl:956
 [5] materialize!
   @ ./broadcast.jl:914
 [6] materialize!
   @ ./broadcast.jl:911
 [7] f
   @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:18

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:1539
  [2] unalias
    @ ./abstractarray.jl:1481 [inlined]
  [3] broadcast_unalias
    @ ./broadcast.jl:977 [inlined]
  [4] preprocess
    @ ./broadcast.jl:984 [inlined]
  [5] preprocess_args
    @ ./broadcast.jl:987 [inlined]
  [6] preprocess_args
    @ ./broadcast.jl:986 [inlined]
  [7] preprocess
    @ ./broadcast.jl:983 [inlined]
  [8] copyto!
    @ ./broadcast.jl:1000 [inlined]
  [9] copyto!
    @ ./broadcast.jl:956 [inlined]
 [10] materialize!
    @ ./broadcast.jl:914 [inlined]
 [11] materialize!
    @ ./broadcast.jl:911 [inlined]
 [12] f
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:18
 [13] ODEFunction
    @ ~/.julia/packages/SciMLBase/fY1XF/src/scimlfunctions.jl:2296 [inlined]
 [14] #224
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:219 [inlined]
 [15] diffejulia__224_11995_inner_1wrap
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:0
 [16] macro expansion
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:6117 [inlined]
 [17] enzyme_call
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5767 [inlined]
 [18] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5644 [inlined]
 [19] autodiff
    @ ~/.julia/packages/Enzyme/wg4P7/src/Enzyme.jl:298 [inlined]
 [20] vec_pjac!(out::ComponentVector{…}, λ::Vector{…}, y::ComponentVector{…}, t::Float64, S::AdjointSensitivityIntegrand{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:294
 [21] AdjointSensitivityIntegrand
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:313 [inlined]
 [22] (::AdjointSensitivityIntegrand{…})(t::Float64)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:325
 [23] evalrule(f::AdjointSensitivityIntegrand{…}, a::Float64, b::Float64, x::Vector{…}, w::Vector{…}, gw::Vector{…}, nrm::typeof(LinearAlgebra.norm))
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/evalrule.jl:0
 [24] #6
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:15 [inlined]
 [25] ntuple
    @ ./ntuple.jl:48 [inlined]
 [26] do_quadgk(f::AdjointSensitivityIntegrand{…}, s::Tuple{…}, n::Int64, atol::Float64, rtol::Float64, maxevals::Int64, nrm::typeof(LinearAlgebra.norm), segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:13
 [27] #50
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:253 [inlined]
 [28] handle_infinities(workfunc::QuadGK.var"#50#51"{…}, f::AdjointSensitivityIntegrand{…}, s::Tuple{…})
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:145
 [29] quadgk(::AdjointSensitivityIntegrand{…}, ::Float64, ::Vararg{…}; atol::Float64, rtol::Float64, maxevals::Int64, order::Int64, norm::Function, segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:252
 [30] _adjoint_sensitivities(sol::ODESolution{…}, sensealg::QuadratureAdjoint{…}, alg::Tsit5{…}; t::Vector{…}, dgdu_discrete::Function, dgdp_discrete::Nothing, dgdu_continuous::Nothing, dgdp_continuous::Nothing, g::Nothing, abstol::Float64, reltol::Float64, callback::Nothing, kwargs::@Kwargs{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:386
 [31] _adjoint_sensitivities
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:328 [inlined]
 [32] #adjoint_sensitivities#63
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/sensitivity_interface.jl:386 [inlined]
 [33] (::SciMLSensitivity.var"#adjoint_sensitivity_backpass#314"{…})(Δ::ODESolution{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/concrete_solve.jl:582
 [34] ZBack
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/chainrules.jl:211 [inlined]
 [35] (::Zygote.var"#291#292"{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:206
 [36] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [37] #solve#51
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1003 [inlined]
 [38] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [39] #291
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:206 [inlined]
 [40] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [41] solve
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:993 [inlined]
 [42] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [43] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:36 [inlined]
 [44] (::Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{Float64, Vector{…}, Tuple{…}}}, Any})(Δ::Float64)
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [45] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{…}}, Any}})(Δ::Float64)
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:91
 [46] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:148
 [47] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:44
Some type information was truncated. Use `show(err)` to see complete types.

ChrisRackauckas commented 3 weeks ago

See if Enzyme.API.runtimeActivity!(true) works?

jClugstor commented 3 weeks ago

Same error

ERROR: Enzyme execution failed.
Mismatched activity for:   %unbox152.fca.1.0.1.insert.pn.extract.0 = phi {} addrspace(10)* [ %unbox152.fca.0.load, %L192 ], [ %getfield, %L175 ] const val:   %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !31, !tbaa !28, !alias.scope !43, !noalias !46, !nonnull !23, !dereferenceable !51, !align !52
 value=Unknown object of type Vector{Float64}
 llvalue=  %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !31, !tbaa !28, !alias.scope !43, !noalias !46, !nonnull !23, !dereferenceable !51, !align !52
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] <
   @ ./int.jl:83
 [2] macro expansion
   @ ./simdloop.jl:72
 [3] copyto!
   @ ./broadcast.jl:1003
 [4] copyto!
   @ ./broadcast.jl:956
 [5] materialize!
   @ ./broadcast.jl:914
 [6] materialize!
   @ ./broadcast.jl:911
 [7] f
   @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:18

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:1539
  [2] unalias
    @ ./abstractarray.jl:1481 [inlined]
  [3] broadcast_unalias
    @ ./broadcast.jl:977 [inlined]
  [4] preprocess
    @ ./broadcast.jl:984 [inlined]
  [5] preprocess_args
    @ ./broadcast.jl:987 [inlined]
  [6] preprocess_args
    @ ./broadcast.jl:986 [inlined]
  [7] preprocess
    @ ./broadcast.jl:983 [inlined]
  [8] copyto!
    @ ./broadcast.jl:1000 [inlined]
  [9] copyto!
    @ ./broadcast.jl:956 [inlined]
 [10] materialize!
    @ ./broadcast.jl:914 [inlined]
 [11] materialize!
    @ ./broadcast.jl:911 [inlined]
 [12] f
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:18
 [13] ODEFunction
    @ ~/.julia/packages/SciMLBase/fY1XF/src/scimlfunctions.jl:2296 [inlined]
 [14] #224
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:219 [inlined]
 [15] diffejulia__224_11995_inner_1wrap
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:0
 [16] macro expansion
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:6117 [inlined]
 [17] enzyme_call
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5767 [inlined]
 [18] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5644 [inlined]
 [19] autodiff
    @ ~/.julia/packages/Enzyme/wg4P7/src/Enzyme.jl:298 [inlined]
 [20] vec_pjac!(out::ComponentVector{…}, λ::Vector{…}, y::ComponentVector{…}, t::Float64, S::AdjointSensitivityIntegrand{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:294
 [21] AdjointSensitivityIntegrand
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:313 [inlined]
 [22] (::AdjointSensitivityIntegrand{…})(t::Float64)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:325
 [23] evalrule(f::AdjointSensitivityIntegrand{…}, a::Float64, b::Float64, x::Vector{…}, w::Vector{…}, gw::Vector{…}, nrm::typeof(LinearAlgebra.norm))
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/evalrule.jl:0
 [24] #6
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:15 [inlined]
 [25] ntuple
    @ ./ntuple.jl:48 [inlined]
 [26] do_quadgk(f::AdjointSensitivityIntegrand{…}, s::Tuple{…}, n::Int64, atol::Float64, rtol::Float64, maxevals::Int64, nrm::typeof(LinearAlgebra.norm), segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:13
 [27] #50
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:253 [inlined]
 [28] handle_infinities(workfunc::QuadGK.var"#50#51"{…}, f::AdjointSensitivityIntegrand{…}, s::Tuple{…})
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:145
 [29] quadgk(::AdjointSensitivityIntegrand{…}, ::Float64, ::Vararg{…}; atol::Float64, rtol::Float64, maxevals::Int64, order::Int64, norm::Function, segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:252
 [30] _adjoint_sensitivities(sol::ODESolution{…}, sensealg::QuadratureAdjoint{…}, alg::Tsit5{…}; t::Vector{…}, dgdu_discrete::Function, dgdp_discrete::Nothing, dgdu_continuous::Nothing, dgdp_continuous::Nothing, g::Nothing, abstol::Float64, reltol::Float64, callback::Nothing, kwargs::@Kwargs{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:386
 [31] _adjoint_sensitivities
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:328 [inlined]
 [32] #adjoint_sensitivities#63
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/sensitivity_interface.jl:386 [inlined]
 [33] (::SciMLSensitivity.var"#adjoint_sensitivity_backpass#314"{…})(Δ::ODESolution{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/concrete_solve.jl:582
 [34] ZBack
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/chainrules.jl:211 [inlined]
 [35] (::Zygote.var"#291#292"{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:206
 [36] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [37] #solve#51
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1003 [inlined]
 [38] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [39] #291
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:206 [inlined]
 [40] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [41] solve
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:993 [inlined]
 [42] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [43] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:38 [inlined]
 [44] (::Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{Float64, Vector{…}, Tuple{…}}}, Any})(Δ::Float64)
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [45] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{…}}, Any}})(Δ::Float64)
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:91
 [46] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:148
 [47] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:46
Some type information was truncated. Use `show(err)` to see complete types.

ChrisRackauckas commented 3 weeks ago

what if you use a let block?

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        let  var"__•1"=var"__•1", __V̇=__V̇
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
    end
end

jClugstor commented 3 weeks ago

Looks like that does it

using Enzyme, PreallocationTools, ComponentArrays, Zygote, OrdinaryDiffEq, SciMLSensitivity
not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        let  var"__•1"=var"__•1", __V̇=__V̇
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                V̇ .= var"•1" .* X
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
    end
end

Enzyme.API.runtimeActivity!(true)

u = ComponentArray(X = [1.0], V = [0.0])
p = ComponentArray(k = 1.0)
f = not_decapode_f()

data_prob = ODEProblem{true, SciMLBase.FullSpecialize}(f, u, (0, 20),p)
sol = solve(data_prob,Tsit5(), saveat = 0.1)
dat = vcat([u.X for u in sol.u]...)

function fake_loss(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, Tsit5(), sensealg = QuadratureAdjoint(autojacvec=EnzymeVJP()))
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss,p) # ((k = -13.667841868057787),)

The Enzyme.API.runtimeActivity!(true) is necessary as well

ChrisRackauckas commented 3 weeks ago

what happens without that?

jClugstor commented 3 weeks ago

without it:

ERROR: Enzyme execution failed.
Mismatched activity for:   %unbox152.fca.1.0.1.insert.pn.extract.0 = phi {} addrspace(10)* [ %unbox152.fca.0.load, %L192 ], [ %getfield, %L175 ] const val:   %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !31, !tbaa !28, !alias.scope !43, !noalias !46, !nonnull !23, !dereferenceable !51, !align !52
 value=Unknown object of type Vector{Float64}
 llvalue=  %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !31, !tbaa !28, !alias.scope !43, !noalias !46, !nonnull !23, !dereferenceable !51, !align !52
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] <
   @ ./int.jl:83
 [2] macro expansion
   @ ./simdloop.jl:72
 [3] copyto!
   @ ./broadcast.jl:1003
 [4] copyto!
   @ ./broadcast.jl:956
 [5] materialize!
   @ ./broadcast.jl:914
 [6] materialize!
   @ ./broadcast.jl:911
 [7] f
   @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:18

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:1539
  [2] unalias
    @ ./abstractarray.jl:1481 [inlined]
  [3] broadcast_unalias
    @ ./broadcast.jl:977 [inlined]
  [4] preprocess
    @ ./broadcast.jl:984 [inlined]
  [5] preprocess_args
    @ ./broadcast.jl:987 [inlined]
  [6] preprocess_args
    @ ./broadcast.jl:986 [inlined]
  [7] preprocess
    @ ./broadcast.jl:983 [inlined]
  [8] copyto!
    @ ./broadcast.jl:1000 [inlined]
  [9] copyto!
    @ ./broadcast.jl:956 [inlined]
 [10] materialize!
    @ ./broadcast.jl:914 [inlined]
 [11] materialize!
    @ ./broadcast.jl:911 [inlined]
 [12] f
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:18
 [13] ODEFunction
    @ ~/.julia/packages/SciMLBase/fY1XF/src/scimlfunctions.jl:2296 [inlined]
 [14] #224
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:219 [inlined]
 [15] diffejulia__224_15724_inner_1wrap
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:0
 [16] macro expansion
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:6117 [inlined]
 [17] enzyme_call
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5767 [inlined]
 [18] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5644 [inlined]
 [19] autodiff
    @ ~/.julia/packages/Enzyme/wg4P7/src/Enzyme.jl:298 [inlined]
 [20] vec_pjac!(out::ComponentVector{…}, λ::Vector{…}, y::ComponentVector{…}, t::Float64, S::AdjointSensitivityIntegrand{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:294
 [21] AdjointSensitivityIntegrand
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:313 [inlined]
 [22] (::AdjointSensitivityIntegrand{…})(t::Float64)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:325
 [23] evalrule(f::AdjointSensitivityIntegrand{…}, a::Float64, b::Float64, x::Vector{…}, w::Vector{…}, gw::Vector{…}, nrm::typeof(LinearAlgebra.norm))
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/evalrule.jl:0
 [24] #6
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:15 [inlined]
 [25] ntuple
    @ ./ntuple.jl:48 [inlined]
 [26] do_quadgk(f::AdjointSensitivityIntegrand{…}, s::Tuple{…}, n::Int64, atol::Float64, rtol::Float64, maxevals::Int64, nrm::typeof(LinearAlgebra.norm), segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:13
 [27] #50
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:253 [inlined]
 [28] handle_infinities(workfunc::QuadGK.var"#50#51"{…}, f::AdjointSensitivityIntegrand{…}, s::Tuple{…})
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:145
 [29] quadgk(::AdjointSensitivityIntegrand{…}, ::Float64, ::Vararg{…}; atol::Float64, rtol::Float64, maxevals::Int64, order::Int64, norm::Function, segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:252
 [30] _adjoint_sensitivities(sol::ODESolution{…}, sensealg::QuadratureAdjoint{…}, alg::Tsit5{…}; t::Vector{…}, dgdu_discrete::Function, dgdp_discrete::Nothing, dgdu_continuous::Nothing, dgdp_continuous::Nothing, g::Nothing, abstol::Float64, reltol::Float64, callback::Nothing, kwargs::@Kwargs{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:386
 [31] _adjoint_sensitivities
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/quadrature_adjoint.jl:328 [inlined]
 [32] #adjoint_sensitivities#63
    @ ~/.julia/packages/SciMLSensitivity/fAsw8/src/sensitivity_interface.jl:386 [inlined]
 [33] (::SciMLSensitivity.var"#adjoint_sensitivity_backpass#314"{…})(Δ::ODESolution{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/fAsw8/src/concrete_solve.jl:582
 [34] ZBack
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/chainrules.jl:211 [inlined]
 [35] (::Zygote.var"#291#292"{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:206
 [36] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [37] #solve#51
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1003 [inlined]
 [38] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [39] #291
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:206 [inlined]
 [40] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [41] solve
    @ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:993 [inlined]
 [42] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [43] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:39 [inlined]
 [44] (::Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{Float64, Vector{…}, Tuple{…}}}, Any})(Δ::Float64)
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface2.jl:0
 [45] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{…}}, Any}})(Δ::Float64)
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:91
 [46] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:148
 [47] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:47
Some type information was truncated. Use `show(err)` to see complete types.

ChrisRackauckas commented 3 weeks ago

Check @code_warntype on that function: what's unstable?

jClugstor commented 3 weeks ago

du = ComponentArray(X = [0.0], V = [0.0])
u = ComponentArray(X = [1.0], V = [3.0])
p = ComponentArray(k = 1.0)

@code_warntype f(du,u,p,1.0)

MethodInstance for (::var"#f#27"{FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}})(::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}}, ::Float64)
  from (::var"#f#27")(du, u, p, t) @ Main ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFix.jl:9
Arguments
  #self#::var"#f#27"{FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}, FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}}
  du::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}
  u::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(X = 1:1, V = 2:2)}}}
  p::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1,)}}}
  t::Float64
Locals
  V̇::Vector{Float64}
  •1::Vector{Float64}
  V::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}
  k::Float64
  X::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}
Body::Nothing
1 ─       (X = Base.getproperty(u, :X))
│         (k = Base.getproperty(p, :k))
│         (V = Base.getproperty(u, :V))
│   %4  = PreallocationTools.get_tmp::Core.Const(PreallocationTools.get_tmp)
│   %5  = Core.getfield(#self#, Symbol("#155#__•1"))::FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}
│         (•1 = (%4)(%5, u))
│   %7  = PreallocationTools.get_tmp::Core.Const(PreallocationTools.get_tmp)
│   %8  = Core.getfield(#self#, Symbol("#156#__V̇"))::FixedSizeDiffCache{Vector{Float64}, Vector{ForwardDiff.Dual{nothing, Float64, 1}}}
│         (V̇ = (%7)(%8, u))
│   %10 = •1::Vector{Float64}
│   %11 = Base.broadcasted(Main.:-, k)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}, Nothing, typeof(-), Tuple{Float64}}
│         Base.materialize!(%10, %11)
│   %13 = V̇::Vector{Float64}
│   %14 = Base.broadcasted(Main.:*, •1, X::Core.PartialStruct(SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Any[Vector{Float64}, Core.Const((1:1,)), Core.Const(0), Core.Const(1)]))::Core.PartialStruct(Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(*), Tuple{Vector{Float64}, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}}}, Any[Core.Const(Base.Broadcast.DefaultArrayStyle{1}()), Core.Const(*), Core.PartialStruct(Tuple{Vector{Float64}, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}}, Any[Vector{Float64}, Core.PartialStruct(SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Any[Vector{Float64}, Core.Const((1:1,)), Core.Const(0), Core.Const(1)])]), Nothing])
│         Base.materialize!(%13, %14)
│   %16 = Main.getproperty(du, :X)::Core.PartialStruct(SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Any[Vector{Float64}, Core.Const((1:1,)), Core.Const(0), Core.Const(1)])
│   %17 = Base.broadcasted(Base.identity, V::Core.PartialStruct(SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Any[Vector{Float64}, Core.Const((2:2,)), Core.Const(1), Core.Const(1)]))::Core.PartialStruct(Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(identity), Tuple{SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}}}, Any[Core.Const(Base.Broadcast.DefaultArrayStyle{1}()), Core.Const(identity), Core.PartialStruct(Tuple{SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}}, Any[Core.PartialStruct(SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Any[Vector{Float64}, Core.Const((2:2,)), Core.Const(1), Core.Const(1)])]), Nothing])
│         Base.materialize!(%16, %17)
│   %19 = Main.getproperty(du, :V)::Core.PartialStruct(SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Any[Vector{Float64}, Core.Const((2:2,)), Core.Const(1), Core.Const(1)])
│   %20 = Base.broadcasted(Base.identity, V̇)::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(identity), Tuple{Vector{Float64}}}
│         Base.materialize!(%19, %20)
└──       return Main.nothing

looks like something in the broadcast of this line V̇ .= var"•1" .* X, there are a couple of Any s in there

wsmoses commented 3 weeks ago

That warning, in the context of a broadcast as is here, isn't an issue of type inference failure, but an inability to ascertain aliasing info, where one of the arguments is non-differentiable

ChrisRackauckas commented 3 weeks ago

So then my understanding is that there's nothing on our end to improve that here, and it just needs an Enzyme improvement? Or should we try something like an in-place map! instead?

ChrisRackauckas commented 2 weeks ago

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        let  var"__•1"=var"__•1", __V̇=__V̇
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                var"•1" .= (.-)(k)
                map!((x,y)->x*y,V̇,var"•1",X)
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
    end
end

@jClugstor check if using a map! like that works. You may also need to do map! on var"•1" .= (.-)(k) and give it another temporary vector so there isn't an aliasing going on. Try that and see if that simplification helps Enzyme.

jClugstor commented 2 weeks ago

with

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        let  var"__•1"=var"__•1", __V̇=__V̇
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)
                V̇ = PreallocationTools.get_tmp(__V̇, u)
                #var"•1" .= (.-)(k)
                map!(x -> .-x, var"•1",k)
                #V̇ .= var"•1" .* X
                map!((x,y)->x*y,V̇,var"•1",X)
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
    end
end

gets the same issue

ERROR: Enzyme execution failed.
Mismatched activity for:   store i64 %11, i64 addrspace(11)* %10, align 8, !dbg !42 const val:   %11 = load i64, i64 addrspace(11)* %9, align 8, !dbg !42
Type tree: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}
 llvalue=  %11 = load i64, i64 addrspace(11)* %9, align 8, !dbg !42
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] _
   @ ~/.julia/packages/SciMLBase/fY1XF/src/problems/ode_problems.jl:118
 [2] ODEProblem
   @ ~/.julia/packages/SciMLBase/fY1XF/src/problems/ode_problems.jl:111
 [3] #remake#653
   @ ~/.julia/packages/SciMLBase/fY1XF/src/remake.jl:133
 [4] remake
   @ ~/.julia/packages/SciMLBase/fY1XF/src/remake.jl:83
 [5] remake
   @ ~/.julia/packages/SciMLBase/fY1XF/src/remake.jl:0

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:1539
  [2] _
    @ ~/.julia/packages/SciMLBase/fY1XF/src/problems/ode_problems.jl:118 [inlined]
  [3] ODEProblem
    @ ~/.julia/packages/SciMLBase/fY1XF/src/problems/ode_problems.jl:111 [inlined]
  [4] #remake#653
    @ ~/.julia/packages/SciMLBase/fY1XF/src/remake.jl:133 [inlined]
  [5] remake
    @ ~/.julia/packages/SciMLBase/fY1XF/src/remake.jl:83 [inlined]
  [6] remake
    @ ~/.julia/packages/SciMLBase/fY1XF/src/remake.jl:0 [inlined]
  [7] augmented_julia_remake_16993_inner_1wrap
    @ ~/.julia/packages/SciMLBase/fY1XF/src/remake.jl:0
  [8] macro expansion
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:6117 [inlined]
  [9] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Val{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Const{…}, ::Const{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5767
 [10] (::Enzyme.Compiler.AugmentedForwardThunk{…})(::Const{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5655
 [11] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(Core.kwcall), df::Nothing, primal_1::@NamedTuple{…}, shadow_1_1::@NamedTuple{…}, primal_2::typeof(remake), shadow_2_1::Nothing, primal_3::ODEProblem{…}, shadow_3_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/wg4P7/src/rules/jitrules.jl:213
 [12] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:47 [inlined]
 [13] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:0 [inlined]
 [14] augmented_julia_fake_loss_16454_inner_1wrap
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:0
 [15] macro expansion
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:6117 [inlined]
 [16] enzyme_call
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5767 [inlined]
 [17] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/wg4P7/src/compiler.jl:5655 [inlined]
 [18] autodiff
    @ ~/.julia/packages/Enzyme/wg4P7/src/Enzyme.jl:242 [inlined]
 [19] autodiff
    @ ~/.julia/packages/Enzyme/wg4P7/src/Enzyme.jl:310 [inlined]
 [20] gradient(rm::ReverseMode{false, FFIABI, false}, f::typeof(fake_loss), x::ComponentVector{Float64, Vector{…}, Tuple{…}})
    @ Enzyme ~/.julia/packages/Enzyme/wg4P7/src/Enzyme.jl:988
 [21] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:56
Some type information was truncated. Use `show(err)` to see complete types.

ChrisRackauckas commented 2 weeks ago

map!(x -> .-x, var"•1",k) That's not correct. It's x - y, there shouldn't be the dot. But x is also the output vector, which map! doesn't allow, so you need another cache vector.

jClugstor commented 2 weeks ago

Doesn't var"•1" .= (.-)(k) just negate every element of k and assign that vector to var"•1" ? Isn't var"•1" the destination, and k the collection to map over?

ChrisRackauckas commented 2 weeks ago

Oh my bad, for some reason I thought it was .-=. Mixed activity stuff landed on Enzyme master, can you try that out?

jClugstor commented 2 weeks ago

Tried Enzyme master out, still an error.

using Enzyme, PreallocationTools, ComponentArrays, Zygote, OrdinaryDiffEq, SciMLSensitivity

Enzyme.API.runtimeActivity!(true)

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        let  var"__•1"=var"__•1", __V̇=__V̇
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)

                V̇ = PreallocationTools.get_tmp(__V̇, u)
                #var"•1" .= (.-)(k)
                map!(x -> .-x, var"•1",k)
                #V̇ .= var"•1" .* X
                map!((x,y)->x*y,V̇,var"•1",X)
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
    end
end

u = ComponentArray(X = [1.0], V = [0.0])
p = ComponentArray(k = 1.0)

du = ComponentArray(X = [0.0], V = [0.0])
u = ComponentArray(X = [1.0], V = [3.0])
p = ComponentArray(k = [1.0])

f = not_decapode_f()

f(du,u,p,1.0)
f(du,u,p,1.0)

data_prob = ODEProblem{true, SciMLBase.FullSpecialize}(f, u, (0, 20),p)
sol = solve(data_prob,Tsit5(), saveat = 0.1)
dat = vcat([u.X for u in sol.u]...)

function fake_loss(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, Tsit5(), sensealg = QuadratureAdjoint(autojacvec=EnzymeVJP()))
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Enzyme.gradient(Reverse, fake_loss, p)

ERROR: Enzyme execution failed.
Mismatched activity for:   store i64 %11, i64 addrspace(11)* %10, align 8, !dbg !42 const val:   %11 = load i64, i64 addrspace(11)* %9, align 8, !dbg !42
Type tree: {[-1]:Pointer, [-1,0]:Pointer, [-1,0,-1]:Float@double, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer, [-1,12]:Integer, [-1,13]:Integer, [-1,14]:Integer, [-1,15]:Integer, [-1,16]:Integer, [-1,17]:Integer, [-1,18]:Integer, [-1,19]:Integer, [-1,20]:Integer, [-1,21]:Integer, [-1,22]:Integer, [-1,23]:Integer, [-1,24]:Integer, [-1,25]:Integer, [-1,26]:Integer, [-1,27]:Integer, [-1,28]:Integer, [-1,29]:Integer, [-1,30]:Integer, [-1,31]:Integer, [-1,32]:Integer, [-1,33]:Integer, [-1,34]:Integer, [-1,35]:Integer, [-1,36]:Integer, [-1,37]:Integer, [-1,38]:Integer, [-1,39]:Integer}
 llvalue=  %11 = load i64, i64 addrspace(11)* %9, align 8, !dbg !42
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] _
   @ ~/.julia/packages/SciMLBase/sakPO/src/problems/ode_problems.jl:118
 [2] ODEProblem
   @ ~/.julia/packages/SciMLBase/sakPO/src/problems/ode_problems.jl:111
 [3] #remake#653
   @ ~/.julia/packages/SciMLBase/sakPO/src/remake.jl:133
 [4] remake
   @ ~/.julia/packages/SciMLBase/sakPO/src/remake.jl:83
 [5] remake
   @ ~/.julia/packages/SciMLBase/sakPO/src/remake.jl:0

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:1612
  [2] _
    @ ~/.julia/packages/SciMLBase/sakPO/src/problems/ode_problems.jl:118 [inlined]
  [3] ODEProblem
    @ ~/.julia/packages/SciMLBase/sakPO/src/problems/ode_problems.jl:111 [inlined]
  [4] #remake#653
    @ ~/.julia/packages/SciMLBase/sakPO/src/remake.jl:133 [inlined]
  [5] remake
    @ ~/.julia/packages/SciMLBase/sakPO/src/remake.jl:83 [inlined]
  [6] remake
    @ ~/.julia/packages/SciMLBase/sakPO/src/remake.jl:0 [inlined]
  [7] augmented_julia_remake_28077_inner_1wrap
    @ ~/.julia/packages/SciMLBase/sakPO/src/remake.jl:0
  [8] macro expansion
    @ ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:6547 [inlined]
  [9] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Val{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Const{…}, ::Const{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:6148
 [10] (::Enzyme.Compiler.AugmentedForwardThunk{…})(::Const{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:6036
 [11] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::typeof(Core.kwcall), df::Nothing, primal_1::@NamedTuple{…}, shadow_1_1::@NamedTuple{…}, primal_2::typeof(remake), shadow_2_1::Nothing, primal_3::ODEProblem{…}, shadow_3_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/VRowe/src/rules/jitrules.jl:311
 [12] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:50 [inlined]
 [13] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:0 [inlined]
 [14] augmented_julia_fake_loss_27525_inner_1wrap
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:0
 [15] macro expansion
    @ ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:6547 [inlined]
 [16] enzyme_call
    @ ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:6148 [inlined]
 [17] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:6036 [inlined]
 [18] autodiff
    @ ~/.julia/packages/Enzyme/VRowe/src/Enzyme.jl:253 [inlined]
 [19] autodiff
    @ ~/.julia/packages/Enzyme/VRowe/src/Enzyme.jl:321 [inlined]
 [20] gradient(rm::ReverseMode{false, FFIABI, false}, f::typeof(fake_loss), x::ComponentVector{Float64, Vector{…}, Tuple{…}})
    @ Enzyme ~/.julia/packages/Enzyme/VRowe/src/Enzyme.jl:1005
 [21] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:59
Some type information was truncated. Use `show(err)` to see complete types.

The error points to the call to new in SciMLBase ode_problems.jl

@add_kwonly function ODEProblem{iip}(f::AbstractODEFunction{iip},
            u0, tspan, p = NullParameters(),
            problem_type = StandardODEProblem();
            kwargs...) where {iip}
        _u0 = prepare_initial_state(u0)
        _tspan = promote_tspan(tspan)
        warn_paramtype(p)
        new{typeof(_u0), typeof(_tspan),
            isinplace(f), typeof(p), typeof(f),
            typeof(kwargs),
            typeof(problem_type)}(f,
            _u0,
            _tspan,
            p,
            kwargs,
            problem_type)
    end

ChrisRackauckas commented 2 weeks ago

No, please just do this one at a time. Just Enzyme in the interior, on the vjp.

jClugstor commented 2 weeks ago

Oh yeah sorry, I should have Zygote there

using Enzyme, PreallocationTools, ComponentArrays, Zygote, OrdinaryDiffEq, SciMLSensitivity

Enzyme.API.runtimeActivity!(false)

not_decapode_f = begin
    function simulate()
        begin
            var"__•1" = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
            __V̇ = PreallocationTools.FixedSizeDiffCache(Vector{Float64}(undef, 1))
        end
        let  var"__•1"=var"__•1", __V̇=__V̇
        f(du, u, p, t) = begin
                begin
                    X = u.X
                    k = p.k
                    V = u.V
                end
                var"•1" = PreallocationTools.get_tmp(var"__•1", u)

                V̇ = PreallocationTools.get_tmp(__V̇, u)
                #var"•1" .= (.-)(k)
                map!(x -> .-x, var"•1",k)
                #V̇ .= var"•1" .* X
                map!((x,y)->x*y,V̇,var"•1",X)
                getproperty(du, :X) .= V
                getproperty(du, :V) .= V̇
                nothing
            end
    end
    end
end

u = ComponentArray(X = [1.0], V = [0.0])
p = ComponentArray(k = 1.0)

du = ComponentArray(X = [0.0], V = [0.0])
u = ComponentArray(X = [1.0], V = [3.0])
p = ComponentArray(k = [1.0])

f = not_decapode_f()

f(du,u,p,1.0)

data_prob = ODEProblem{true, SciMLBase.FullSpecialize}(f, u, (0, 20),p)
sol = solve(data_prob,Tsit5(), saveat = 0.1)
dat = vcat([u.X for u in sol.u]...)

function fake_loss(p)
    prob = remake(data_prob, p = p)
    soln = solve(prob, Tsit5(), sensealg = QuadratureAdjoint(autojacvec=EnzymeVJP()))
    @info("Done")

    # return soln(tₑ)
    sum(last(soln)) # last, not soln(tₑ) because to avoid interpolation fails when AD fails.
end

Zygote.gradient(fake_loss,p)

ERROR: Enzyme execution failed.
Mismatched activity for:   %.pn249 = phi {} addrspace(10)* [ %getfield, %L325 ], [ %43, %L338 ] const val:   %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !28, !tbaa !25, !alias.scope !40, !noalias !43, !nonnull !20, !dereferenceable !48, !align !49
 value=Unknown object of type Vector{Float64}
 llvalue=  %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !28, !tbaa !25, !alias.scope !40, !noalias !43, !nonnull !20, !dereferenceable !48, !align !49
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] copyto!
   @ ./abstractarray.jl:1068
 [2] copyto!
   @ ./broadcast.jl:997
 [3] copyto!
   @ ./broadcast.jl:956
 [4] materialize!
   @ ./broadcast.jl:914
 [5] materialize!
   @ ./broadcast.jl:911
 [6] f
   @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:25

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:1612
  [2] unalias
    @ ./abstractarray.jl:1481 [inlined]
  [3] copyto!
    @ ./abstractarray.jl:1067 [inlined]
  [4] copyto!
    @ ./broadcast.jl:997 [inlined]
  [5] copyto!
    @ ./broadcast.jl:956 [inlined]
  [6] materialize!
    @ ./broadcast.jl:914 [inlined]
  [7] materialize!
    @ ./broadcast.jl:911 [inlined]
  [8] f
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:25
  [9] ODEFunction
    @ ~/.julia/packages/SciMLBase/sakPO/src/scimlfunctions.jl:2296 [inlined]
 [10] #224
    @ ~/.julia/packages/SciMLSensitivity/waEMv/src/quadrature_adjoint.jl:219 [inlined]
 [11] diffejulia__224_14657_inner_1wrap
    @ ~/.julia/packages/SciMLSensitivity/waEMv/src/quadrature_adjoint.jl:0
 [12] macro expansion
    @ ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:6547 [inlined]
 [13] enzyme_call
    @ ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:6148 [inlined]
 [14] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/VRowe/src/compiler.jl:6025 [inlined]
 [15] autodiff
    @ ~/.julia/packages/Enzyme/VRowe/src/Enzyme.jl:309 [inlined]
 [16] vec_pjac!(out::ComponentVector{…}, λ::Vector{…}, y::ComponentVector{…}, t::Float64, S::AdjointSensitivityIntegrand{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/waEMv/src/quadrature_adjoint.jl:294
 [17] AdjointSensitivityIntegrand
    @ ~/.julia/packages/SciMLSensitivity/waEMv/src/quadrature_adjoint.jl:313 [inlined]
 [18] (::AdjointSensitivityIntegrand{…})(t::Float64)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/waEMv/src/quadrature_adjoint.jl:325
 [19] evalrule(f::AdjointSensitivityIntegrand{…}, a::Float64, b::Float64, x::Vector{…}, w::Vector{…}, gw::Vector{…}, nrm::typeof(LinearAlgebra.norm))
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/evalrule.jl:0
 [20] #6
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:15 [inlined]
 [21] ntuple
    @ ./ntuple.jl:48 [inlined]
 [22] do_quadgk(f::AdjointSensitivityIntegrand{…}, s::Tuple{…}, n::Int64, atol::Float64, rtol::Float64, maxevals::Int64, nrm::typeof(LinearAlgebra.norm), segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:13
 [23] #50
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:253 [inlined]
 [24] handle_infinities(workfunc::QuadGK.var"#50#51"{…}, f::AdjointSensitivityIntegrand{…}, s::Tuple{…})
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:145
 [25] quadgk(::AdjointSensitivityIntegrand{…}, ::Float64, ::Vararg{…}; atol::Float64, rtol::Float64, maxevals::Int64, order::Int64, norm::Function, segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:252
 [26] _adjoint_sensitivities(sol::ODESolution{…}, sensealg::QuadratureAdjoint{…}, alg::Tsit5{…}; t::Vector{…}, dgdu_discrete::Function, dgdp_discrete::Nothing, dgdu_continuous::Nothing, dgdp_continuous::Nothing, g::Nothing, abstol::Float64, reltol::Float64, callback::Nothing, kwargs::@Kwargs{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/waEMv/src/quadrature_adjoint.jl:386
 [27] _adjoint_sensitivities
    @ ~/.julia/packages/SciMLSensitivity/waEMv/src/quadrature_adjoint.jl:328 [inlined]
 [28] #adjoint_sensitivities#63
    @ ~/.julia/packages/SciMLSensitivity/waEMv/src/sensitivity_interface.jl:386 [inlined]
 [29] (::SciMLSensitivity.var"#adjoint_sensitivity_backpass#314"{…})(Δ::ODESolution{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/waEMv/src/concrete_solve.jl:582
 [30] ZBack
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/chainrules.jl:211 [inlined]
 [31] (::Zygote.var"#291#292"{…})(Δ::ODESolution{…})
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:206
 [32] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [33] #solve#51
    @ ~/.julia/packages/DiffEqBase/DS1sd/src/solve.jl:1003 [inlined]
 [34] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:-37
 [35] #291
    @ ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/lib/lib.jl:206 [inlined]
 [36] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [37] solve
    @ ~/.julia/packages/DiffEqBase/DS1sd/src/solve.jl:993 [inlined]
 [38] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:-37
 [39] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:51 [inlined]
 [40] (::Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{Float64, Vector{…}, Tuple{…}}}, Any})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:-37
 [41] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{…}}, Any}})(Δ::Float64)
    @ Zygote ~/Documents/Work/dev/DecapodesAutoDiff/Zygote.jl/src/compiler/interface.jl:91
 [42] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1:1,)}}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:148
 [43] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:58
Some type information was truncated. Use `show(err)` to see complete types.

It's pointing to the line getproperty(du, :X) .= V

wsmoses commented 2 weeks ago

yeah Enzyme (reasonably) thinks that those two could alias with each other.

jClugstor commented 2 weeks ago

Sorry, what two things does it think could alias? du.X and V ?

jClugstor commented 3 days ago

I tried replacing the getproperty calls to setproperty!

 #getproperty(du, :X) .= V
                #getproperty(du, :V) .= V̇
                setproperty!(du,:X, V)
                setproperty!(du,:V, V̇)

And it still errors with mismatched activity:

ERROR: Enzyme execution failed.
Mismatched activity for:   %unbox27.fca.1.0.1.insert.pn.extract.0 = phi {} addrspace(10)* [ %unbox27.fca.0.load, %L29 ], [ %getfield4, %L31 ] const val:   %getfield4 = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr3 unordered, align 8, !dbg !114, !tbaa !25, !alias.scope !33, !noalias !36, !nonnull !19, !dereferenceable !41, !align !42
 value=Unknown object of type Vector{Float64}
 llvalue=  %getfield4 = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr3 unordered, align 8, !dbg !114, !tbaa !25, !alias.scope !33, !noalias !36, !nonnull !19, !dereferenceable !41, !align !42
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] setindex!
   @ ./array.jl:1021
 [2] setindex!
   @ ./array.jl:1041
 [3] macro expansion
   @ ~/.julia/packages/ComponentArrays/joxVV/src/array_interface.jl:0
 [4] _setindex!
   @ ~/.julia/packages/ComponentArrays/joxVV/src/array_interface.jl:131

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/YQwVA/src/compiler.jl:1620
  [2] checkbounds
    @ ./abstractarray.jl:702 [inlined]
  [3] getindex
    @ ./subarray.jl:322 [inlined]
  [4] setindex!
    @ ./array.jl:1040 [inlined]
  [5] macro expansion
    @ ~/.julia/packages/ComponentArrays/joxVV/src/array_interface.jl:0 [inlined]
  [6] _setindex!
    @ ~/.julia/packages/ComponentArrays/joxVV/src/array_interface.jl:131
  [7] setproperty!
    @ ~/.julia/packages/ComponentArrays/joxVV/src/namedtuple_interface.jl:17 [inlined]
  [8] f
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:26
  [9] ODEFunction
    @ ~/.julia/packages/SciMLBase/sakPO/src/scimlfunctions.jl:2296 [inlined]
 [10] #224
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:222 [inlined]
 [11] diffejulia__224_11092_inner_1wrap
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:0
 [12] macro expansion
    @ ~/.julia/packages/Enzyme/YQwVA/src/compiler.jl:6606 [inlined]
 [13] enzyme_call
    @ ~/.julia/packages/Enzyme/YQwVA/src/compiler.jl:6207 [inlined]
 [14] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/YQwVA/src/compiler.jl:6084 [inlined]
 [15] autodiff
    @ ~/.julia/packages/Enzyme/YQwVA/src/Enzyme.jl:309 [inlined]
 [16] vec_pjac!(out::ComponentVector{…}, λ::Vector{…}, y::ComponentVector{…}, t::Float64, S::AdjointSensitivityIntegrand{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:297
 [17] AdjointSensitivityIntegrand
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:316 [inlined]
 [18] (::AdjointSensitivityIntegrand{…})(t::Float64)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:328
 [19] evalrule(f::AdjointSensitivityIntegrand{…}, a::Float64, b::Float64, x::Vector{…}, w::Vector{…}, gw::Vector{…}, nrm::typeof(LinearAlgebra.norm))
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/evalrule.jl:0
 [20] #6
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:15 [inlined]
 [21] ntuple
    @ ./ntuple.jl:48 [inlined]
 [22] do_quadgk(f::AdjointSensitivityIntegrand{…}, s::Tuple{…}, n::Int64, atol::Float64, rtol::Float64, maxevals::Int64, nrm::typeof(LinearAlgebra.norm), segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:13
 [23] #50
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:253 [inlined]
 [24] handle_infinities(workfunc::QuadGK.var"#50#51"{…}, f::AdjointSensitivityIntegrand{…}, s::Tuple{…})
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:145
 [25] quadgk(::AdjointSensitivityIntegrand{…}, ::Float64, ::Vararg{…}; atol::Float64, rtol::Float64, maxevals::Int64, order::Int64, norm::Function, segbuf::Nothing)
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:252
 [26] _adjoint_sensitivities(sol::ODESolution{…}, sensealg::QuadratureAdjoint{…}, alg::Tsit5{…}; t::Vector{…}, dgdu_discrete::Function, dgdp_discrete::Nothing, dgdu_continuous::Nothing, dgdp_continuous::Nothing, g::Nothing, abstol::Float64, reltol::Float64, callback::Nothing, kwargs::@Kwargs{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:389
 [27] _adjoint_sensitivities
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:331 [inlined]
 [28] #adjoint_sensitivities#63
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/sensitivity_interface.jl:401 [inlined]
 [29] (::SciMLSensitivity.var"#adjoint_sensitivity_backpass#308"{…})(Δ::ODESolution{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/4YtYh/src/concrete_solve.jl:625
 [30] ZBack
    @ ~/.julia/packages/Zygote/nsBv0/src/compiler/chainrules.jl:211 [inlined]
 [31] (::Zygote.var"#291#292"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:206
 [32] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [33] #solve#51
    @ ~/.julia/packages/DiffEqBase/c8MAQ/src/solve.jl:1003 [inlined]
 [34] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [35] #291
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:206 [inlined]
 [36] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [37] solve
    @ ~/.julia/packages/DiffEqBase/c8MAQ/src/solve.jl:993 [inlined]
 [38] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [39] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:52 [inlined]
 [40] (::Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{Float64, Vector{…}, Tuple{…}}}, Any})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [41] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{…}}, Any}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:91
 [42] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1:1,)}}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:148
 [43] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:59
Some type information was truncated. Use `show(err)` to see complete types.

Could this be because of something in ComponentArrays?

ChrisRackauckas commented 3 days ago

Change those broadcasts to copyto!s?