Harmonic Oscillator Decapode Calibration

[jClugstor]

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]

Still gives mismatched activity

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)
                copyto!(getproperty(du, :X),V)
                copyto!(getproperty(du, :V),V̇)
                #setproperty!(du,:X, V)
                #setproperty!(du,:V, V̇)
                nothing
            end
    end
    end
end

ERROR: Enzyme execution failed.
Mismatched activity for:   %.pn214 = phi {} addrspace(10)* [ %getfield, %L325 ], [ %43, %L332 ] const val:   %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !25, !tbaa !22, !alias.scope !37, !noalias !40, !nonnull !17, !dereferenceable !45, !align !46
 value=Unknown object of type Vector{Float64}
 llvalue=  %getfield = load atomic {} addrspace(10)*, {} addrspace(10)* addrspace(11)* %getfield_addr unordered, align 8, !dbg !25, !tbaa !22, !alias.scope !37, !noalias !40, !nonnull !17, !dereferenceable !45, !align !46
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] f
   @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:24

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/YQwVA/src/compiler.jl:1620
  [2] unalias
    @ ./abstractarray.jl:1481 [inlined]
  [3] copyto!
    @ ./abstractarray.jl:1067 [inlined]
  [4] f
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:24
  [5] ODEFunction
    @ ~/.julia/packages/SciMLBase/sakPO/src/scimlfunctions.jl:2296 [inlined]
  [6] #224
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:222 [inlined]
  [7] diffejulia__224_8189_inner_1wrap
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:0
  [8] macro expansion
    @ ~/.julia/packages/Enzyme/YQwVA/src/compiler.jl:6606 [inlined]
  [9] enzyme_call
    @ ~/.julia/packages/Enzyme/YQwVA/src/compiler.jl:6207 [inlined]
 [10] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/YQwVA/src/compiler.jl:6084 [inlined]
 [11] autodiff
    @ ~/.julia/packages/Enzyme/YQwVA/src/Enzyme.jl:309 [inlined]
 [12] vec_pjac!(out::ComponentVector{…}, λ::Vector{…}, y::ComponentVector{…}, t::Float64, S::AdjointSensitivityIntegrand{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:297
 [13] AdjointSensitivityIntegrand
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:316 [inlined]
 [14] (::AdjointSensitivityIntegrand{…})(t::Float64)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:328
 [15] 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
 [16] #6
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:15 [inlined]
 [17] ntuple
    @ ./ntuple.jl:48 [inlined]
 [18] 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
 [19] #50
    @ ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:253 [inlined]
 [20] handle_infinities(workfunc::QuadGK.var"#50#51"{…}, f::AdjointSensitivityIntegrand{…}, s::Tuple{…})
    @ QuadGK ~/.julia/packages/QuadGK/OtnWt/src/adapt.jl:145
 [21] 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
 [22] _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
 [23] _adjoint_sensitivities
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/quadrature_adjoint.jl:331 [inlined]
 [24] #adjoint_sensitivities#63
    @ ~/.julia/packages/SciMLSensitivity/4YtYh/src/sensitivity_interface.jl:401 [inlined]
 [25] (::SciMLSensitivity.var"#adjoint_sensitivity_backpass#308"{…})(Δ::ODESolution{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/4YtYh/src/concrete_solve.jl:625
 [26] ZBack
    @ ~/.julia/packages/Zygote/nsBv0/src/compiler/chainrules.jl:211 [inlined]
 [27] (::Zygote.var"#291#292"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:206
 [28] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [29] #solve#51
    @ ~/.julia/packages/DiffEqBase/c8MAQ/src/solve.jl:1003 [inlined]
 [30] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [31] #291
    @ ~/.julia/packages/Zygote/nsBv0/src/lib/lib.jl:206 [inlined]
 [32] (::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [33] solve
    @ ~/.julia/packages/DiffEqBase/c8MAQ/src/solve.jl:993 [inlined]
 [34] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [35] fake_loss
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:52 [inlined]
 [36] (::Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{Float64, Vector{…}, Tuple{…}}}, Any})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [37] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{typeof(fake_loss), ComponentVector{…}}, Any}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:91
 [38] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(k = 1:1,)}}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:148
 [39] top-level scope
    @ ~/Documents/Work/dev/DecapodesAutoDiff/EnzymeFix/EnzymeFixMaps.jl:59
Some type information was truncated. Use `show(err)` to see complete types.

[ChrisRackauckas]

@wsmoses I don't quite understand this one.

[wsmoses]

Hm is this not an error saying to use runtime activity ? Due to some aliasing issue in the copy?

[ChrisRackauckas]

None of the arrays can be aliasing though, they are all given as different buffers.

[wsmoses]

Sure, but Julia doesn't know that

[wsmoses]

In particular in line 24 of your snippet it fails to statically prove this check won't be needed: https://github.com/JuliaLang/julia/blob/4954197196d657d14edd3e9c61ac101866e6fa25/base/abstractarray.jl#L1067

[wsmoses]

@ChrisRackauckas not sure on the semantics of preallocationtools, but yo may be able to give it an Enzyme.EnzymeRules.noalias attribute [to specify that the return of PreallocationTools.get_tmp will never alias with any other memory buffer

[ChrisRackauckas]

There's no docs on the output and no test that uses it? I presume it outputs a bool?

@jClugstor can you try adding to the script:

EnzymeCore.EnzymeRules.noalias(::Typeof(PreallocationTools.get_tmp), args...) =  true

[wsmoses]

Yeah it’s super experimental rn but yeah. It’s used in cuda.jl rn if you want to take a look (who is its only user presently).

I am worried about the semantic mismatch here tho so I want to understand the guarantees provided by preallocationtools

[ChrisRackauckas]

https://github.com/SciML/PreallocationTools.jl/blob/master/src/PreallocationTools.jl#L13

It always builds its own cache buffers and then reinterprets those as needed for different types. So the cache construction ensures that the only way to get the same memory would be to get_tmp from the same buffer, which would just give incorrect results and so that wouldn't be a thing a user could actually do in a real code.