Hi!

In my original problem the computation of the objective value function is expensive and during the computation I can collaterally compute the constraints of my problem. To make my problem more efficient I tried to implement a version where I store the constraint info associated with the last evaluated solution in a Cache. When trying to use AutoEnzyme() in Optimization.jl it seems to break because Im forced to pass the Cache through the parameter vector p and my understanding is that this is treated as a Const by Enzyme under the hood. If I explicitly use autodiff on a version of my objective function where the Cache is a separate argument and passing the Cache as DuplicatedNoNeed it will differentiate the function correctly. To give a better sense of what I’m trying here is my MWE:

using Optimization 
using OptimizationMOI 
using Ipopt 
using Enzyme 
using SpecialFunctions

cdf_eval(x, mean, stddev) =  0.5 * (1 + erf((x - mean) / (stddev * sqrt(2.0))))

struct Student{T}
    θ::T
    σ::T 
    u1::T
    u2::T  
end

mutable struct Cache{T}
    last_cutoffs::Vector{T}          
    welfare::T                 
    demand::Vector{T}          
end

function student_choice!(cutoffs, student, prob)
    """
    simple function to simulate the school choice a student given a cutoff vector 
    and compute its expected value. function modified prob in-place for efficiency motives 
    """
    # unpack student 
    (; θ, σ, u1, u2) = student 
    # we constraint c2 > c1 
    if u2 ≥ u1
        prob[2] = 1.0 - cdf_eval(cutoffs[2], θ, σ)
        prob[1] = prob[2] - cdf_eval(cutoffs[1], θ, σ)
        expected_u = prob[2] * u2 + prob[1] * u1
    else
        prob[2] = 0.0
        prob[1] = 1.0 - cdf_eval(cutoffs[1], θ, σ)
        expected_u = prob[1] * u1
    end
    return expected_u
end

function eval!(cutoffs, p)
    """
    objetive value function that updates the Cache with demand calculations for constraint and also welfare calculation for 
    objective value 
    """
    array_students = p[1]::Vector{Student{Float64}}
    cache = p[2]::Cache
    demand = cache.demand
    demand .= zero(eltype(cutoffs))

    probs = zeros(eltype(cutoffs), 2)
    welfare = zero(eltype(cutoffs))
    @inbounds for student ∈ array_students
        probs .= zero(eltype(cutoffs))
        welfare += student_choice!(cutoffs, student, probs)
        demand .+= probs
    end

    cache.last_cutoffs .= cutoffs
    cache.welfare = welfare
    return welfare             
end

function f(cutoffs, p)
    last = cache.last_cutoffs
    last === cutoffs || eval!(cutoffs, p)
    return cache.welfare
end

function cons(res, cutoffs, p)
    cache = p[2]::Cache
    capacities = p[3]::Vector{Float64}
    last = cache.last_cutoffs
    last === cutoffs || eval!(cutoffs, p)
    res[1] = capacities[1] - cache.demand[1] 
    res[2] = capacities[2] - cache.demand[2] 
    res[3] = cutoffs[2] - cutoffs[1]
    return nothing
end
function main()
    ## SIMULATE DATA ## 
    number_students = 1000
    array_students = Array{Student{Float64}}(undef, number_students)
    for student_id ∈ 1:1000
        θ = rand()
        σ = 0.2
        u1 = rand() + 1.0 
        u2 = rand() + 1.0
        array_students[student_id] = Student(θ, σ, u1, u2)
    end
    capacities = [400.0, 200.0]
    ## PERFORM OPTIMIZATION ## 
    cache = Cache(zeros(2), 0.0, zeros(2))
    p = [array_students, cache, capacities]
    x0 = rand(2)
    optprob = OptimizationFunction(f, AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse)), cons = cons)
    prob = OptimizationProblem(optprob, x0, p, lcons = [0.0, 0.0, 0.0], ucons = [Inf, Inf, Inf])
    sol = solve(prob, Ipopt.Optimizer())
end
main()

If there a way around this? Like telling AutoEnzyme that a given element of the parameter vector needs to be treated as differentiable even though we are not optimizing over it?

The error when I try the above is:

Allocation could not have its type statically determined   %150 = call noalias nonnull "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @julia.gc_alloc_obj({}* nonnull %6, i64 %143, {} addrspace(10)* nonnull %148) #42, !dbg !118
Stacktrace:
  [1] shadow_alloc_rewrite(V::Ptr{…}, gutils::Ptr{…}, Orig::Ptr{…}, idx::UInt64, prev::Ptr{…}, used::UInt8)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:681
  [2] EnzymeCreateForwardDiff(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnValue::Bool, mode::Enzyme.API.CDerivativeMode, runtimeActivity::Bool, width::Int64, additionalArg::Ptr{…}, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…})
    @ Enzyme.API ~/.julia/packages/Enzyme/hu9gq/src/api.jl:338
  [3] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{…} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:1793
  [4] codegen(output::Symbol, job::GPUCompiler.CompilerJob{…}; libraries::Bool, deferred_codegen::Bool, optimize::Bool, toplevel::Bool, strip::Bool, validate::Bool, only_entry::Bool, parent_job::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:4669
  [5] codegen
    @ ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:3455 [inlined]
  [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget, Enzyme.Compiler.EnzymeCompilerParams}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:5533
  [7] _thunk
    @ ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:5533 [inlined]
  [8] cached_compilation
    @ ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:5585 [inlined]
  [9] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:5696
 [10] thunk_generator(world::UInt64, source::LineNumberNode, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:5881
 [11] autodiff(::ForwardMode{…}, ::Const{…}, ::Type{…}, ::Const{…}, ::BatchDuplicated{…}, ::BatchDuplicatedNoNeed{…}, ::Const{…}, ::Const{…}, ::Const{…}, ::Const{…}, ::Const{…}, ::Const{…})
    @ Enzyme ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:641
 [12] autodiff
    @ ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:545 [inlined]
 [13] autodiff
    @ ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:517 [inlined]
 [14] (::OptimizationEnzymeExt.var"#lag_h!#54"{…})(h::SubArray{…}, θ::Vector{…}, σ::Float64, μ::Vector{…}, p::Vector{…})
    @ OptimizationEnzymeExt ~/.julia/packages/OptimizationBase/UXLhR/ext/OptimizationEnzymeExt.jl:344
 [15] lag_h!
    @ ~/.julia/packages/OptimizationBase/UXLhR/ext/OptimizationEnzymeExt.jl:341 [inlined]
 [16] eval_hessian_lagrangian(evaluator::OptimizationMOI.MOIOptimizationNLPEvaluator{…}, h::SubArray{…}, x::Vector{…}, σ::Float64, μ::SubArray{…})
    @ OptimizationMOI ~/.julia/packages/OptimizationMOI/L0B28/src/nlp.jl:378
 [17] eval_hessian_lagrangian(model::IpoptMathOptInterfaceExt.Optimizer, H::Vector{Float64}, x::Vector{Float64}, σ::Float64, μ::Vector{Float64})
    @ IpoptMathOptInterfaceExt ~/.julia/packages/Ipopt/Fbuwv/ext/IpoptMathOptInterfaceExt/MOI_wrapper.jl:1303
 [18] (::IpoptMathOptInterfaceExt.var"#eval_h_cb#11"{…})(x::Vector{…}, rows::Vector{…}, cols::Vector{…}, obj_factor::Float64, lambda::Vector{…}, values::Vector{…})
    @ IpoptMathOptInterfaceExt ~/.julia/packages/Ipopt/Fbuwv/ext/IpoptMathOptInterfaceExt/MOI_wrapper.jl:1387
 [19] _Eval_H_CB(n::Int32, x_ptr::Ptr{…}, ::Int32, obj_factor::Float64, m::Int32, lambda_ptr::Ptr{…}, ::Int32, nele_hess::Int32, iRow::Ptr{…}, jCol::Ptr{…}, values_ptr::Ptr{…}, user_data::Ptr{…})
    @ Ipopt ~/.julia/packages/Ipopt/Fbuwv/src/C_wrapper.jl:0
 [20] #5
    @ ~/.julia/packages/Ipopt/Fbuwv/src/C_wrapper.jl:407 [inlined]
 [21] disable_sigint(f::Ipopt.var"#5#6"{IpoptProblem, Base.RefValue{Float64}})
    @ Base ./c.jl:167
 [22] IpoptSolve
    @ ~/.julia/packages/Ipopt/Fbuwv/src/C_wrapper.jl:406 [inlined]
 [23] optimize!(model::IpoptMathOptInterfaceExt.Optimizer)
    @ IpoptMathOptInterfaceExt ~/.julia/packages/Ipopt/Fbuwv/ext/IpoptMathOptInterfaceExt/MOI_wrapper.jl:1523
 [24] __solve(cache::OptimizationMOI.MOIOptimizationNLPCache{OptimizationMOI.MOIOptimizationNLPEvaluator{…}, IpoptMathOptInterfaceExt.Optimizer})
    @ OptimizationMOI ~/.julia/packages/OptimizationMOI/L0B28/src/nlp.jl:552
 [25] solve!(cache::OptimizationMOI.MOIOptimizationNLPCache{OptimizationMOI.MOIOptimizationNLPEvaluator{…}, IpoptMathOptInterfaceExt.Optimizer})
    @ SciMLBase ~/.julia/packages/SciMLBase/iHgIu/src/solve.jl:227
 [26] solve(::OptimizationProblem{…}, ::IpoptMathOptInterfaceExt.Optimizer; kwargs::@Kwargs{})
    @ SciMLBase ~/.julia/packages/SciMLBase/iHgIu/src/solve.jl:129
 [27] solve(::OptimizationProblem{…}, ::IpoptMathOptInterfaceExt.Optimizer)
    @ SciMLBase ~/.julia/packages/SciMLBase/iHgIu/src/solve.jl:126
 [28] main()
    @ Main ./REPL[37]:19
 [29] top-level scope
    @ REPL[40]:1
Some type information was truncated. Use `show(err)` to see complete types.

Nowadays, DifferentiationInterface.jl supports both Constants and Caches, but the Optimization.jl bindings consider everything to be constant. Perhaps it would be worth splitting p in two?

1 Like

Hi there, thanks for the comment. I do not really 100% follow your suggestion, can you elaborate a bit on what do you mean by splitting p into two and how that would allow me to tell Optimization.jl that one part is differentiable?

Basically if there was a p_constant and a p_cache this would make it easier to provide the correct annotations to Enzyme and other backends

Oh I thought you were suggesting that to me as an immediate solution to my problem, but I am guessing your question is addressed to the Optimization.jl developers.

1 Like

Yes indeed, sorry for the confusion.

To handle your immediate issue, I guess you could provide gradients to Optimization.jl manually, leveraging DI’s Cache context?

Enzyme Dev here.

Alternate suggestion, can you make the original function a closure, and mark the AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=DuplicatedNoNeed).

This might require a quick patch to optimization which I can also help with if it doesn’t work first go.

I think this is what you want from the first comment, though note that Cache is not an Enzyme or Optimization primitive, and DI does not have support for DuplicatedNoNeed.

1 Like

Indeed, at the moment OptimizationBase doesn’t call Enzyme through DI so it will happily ignore both the mode and the function_annotation argument (IIUC).

Right now, if you select function_annotation=DuplicatedNoNeed and call Enzyme through DI, it will just fall back on Duplicated (which is not wrong but slightly less efficient). I have opened a PR to correct this behavior:

Hi there William, thanks for your comment. This is how I implemented the closures:

cdf_eval(x, mean, stddev) =  0.5 * (1 + erf((x - mean) / (stddev * sqrt(2.0))))

struct Student{T}
    θ::T
    σ::T 
    u1::T
    u2::T  
end

mutable struct Cache{T}
    last_cutoffs::Vector{T}          
    welfare::T                 
    demand::Vector{T}          
end

function student_choice!(cutoffs, student, prob)
    """
    simple function to simulate the school choice a student given a cutoff vector 
    and compute its expected value. function modified prob in-place for efficiency motives 
    """
    # unpack student 
    (; θ, σ, u1, u2) = student 
    # we constraint c2 > c1 
    if u2 ≥ u1
        prob[2] = 1.0 - cdf_eval(cutoffs[2], θ, σ)
        prob[1] = prob[2] - cdf_eval(cutoffs[1], θ, σ)
        expected_u = prob[2] * u2 + prob[1] * u1
    else
        prob[2] = 0.0
        prob[1] = 1.0 - cdf_eval(cutoffs[1], θ, σ)
        expected_u = prob[1] * u1
    end
    return expected_u
end

function eval!(cutoffs, cache, p)
    array_students = p[1]::Vector{Student{Float64}}
    demand = cache.demand
    demand .= zero(eltype(cutoffs))

    probs = zeros(eltype(cutoffs), 2)
    welfare = zero(eltype(cutoffs))
    @inbounds for student ∈ array_students
        probs .= zero(eltype(cutoffs))
        welfare += student_choice!(cutoffs, student, probs)
        demand .+= probs
    end

    cache.last_cutoffs .= cutoffs
    cache.welfare = welfare
    return welfare             
end

function f(cutoffs, cache, p)
    last = cache.last_cutoffs
    last === cutoffs || eval!(cutoffs, cache, p)
    return cache.welfare
end

# closure 

function g(cache)
    f_closure = let cache = cache
        (cutoffs, p) -> f(cutoffs, cache, p)
    end
    return f_closure
end

function cons(res, cutoffs, cache, p)
    capacities = p[2]::Vector{Float64}
    last = cache.last_cutoffs
    last === cutoffs || eval!(cutoffs, cache, p)
    res[1] = capacities[1] - cache.demand[1] 
    res[2] = capacities[2] - cache.demand[2] 
    res[3] = cutoffs[2] - cutoffs[1]
    return nothing
end

function g_cons(cache)
    cons_closure = let cache = cache 
        (res, cutoffs, p) -> cons(res, cutoffs, cache, p)
    end 
    return cons_closure
end
function main()
    ## SIMULATE DATA ## 
    number_students = 1000
    array_students = Array{Student{Float64}}(undef, number_students)
    for student_id ∈ 1:1000
        θ = rand()
        σ = 0.2
        u1 = rand() + 1.0 
        u2 = rand() + 1.0
        array_students[student_id] = Student(θ, σ, u1, u2)
    end
    capacities = [400.0, 200.0]
    ## PERFORM OPTIMIZATION ## 
    #cache = Cache(zeros(2), 0.0, zeros(2))
    #p = [array_students, cache, capacities]
    #x0 = rand(2)
    #optprob = OptimizationFunction(f, AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse)), cons = cons)
    #prob = OptimizationProblem(optprob, x0, p, lcons = [0.0, 0.0, 0.0], ucons = [Inf, Inf, Inf])
    #sol = solve(prob, Ipopt.Optimizer())

    ## USE CLOSURE ## 
    cache = Cache(zeros(2), 0.0, zeros(2))
    p = [array_students, capacities]
    x0 = rand(2)
    f_closure = g(cache)
    cons_closure = g_cons(cache)
    optprob = OptimizationFunction(f_closure, AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse), function_annotation=DuplicatedNoNeed), cons = cons_closure)
    prob = OptimizationProblem(optprob, x0, p, lcons = [0.0, 0.0, 0.0], ucons = [Inf, Inf, Inf])
    sol = solve(prob, Ipopt.Optimizer())
end
main()

But it is still failing with the error:

******************************************************************************
This program contains Ipopt, a library for large-scale nonlinear optimization.
 Ipopt is released as open source code under the Eclipse Public License (EPL).
         For more information visit https://github.com/coin-or/Ipopt
******************************************************************************

This is Ipopt version 3.14.17, running with linear solver MUMPS 5.7.3.

Number of nonzeros in equality constraint Jacobian...:        0
Number of nonzeros in inequality constraint Jacobian.:        6
Number of nonzeros in Lagrangian Hessian.............:        3

ERROR: 
If you are using Enzyme by selecting the `AutoEnzyme` object from ADTypes, you may want to try setting the `function_annotation` option as follows:

        AutoEnzyme(; function_annotation=Enzyme.Duplicated)

This hint appears because DifferentiationInterface and Enzyme are both loaded. It does not necessarily imply that Enzyme is being called through DifferentiationInterface.

Function argument passed to autodiff cannot be proven readonly.
If the the function argument cannot contain derivative data, instead call autodiff(Mode, Const(f), ...)
See https://enzyme.mit.edu/index.fcgi/julia/stable/faq/#Activity-of-temporary-storage for more information.
The potentially writing call is   call void @llvm.memset.p13i8.i64(i8 addrspace(13)* align 8 %arrayptr11.pre238283, i8 noundef 0, i64 %118, i1 noundef false), !dbg !330, !tbaa !87, !alias.scope !90, !noalias !91, using   %arrayptr11.pre238283 = load i8 addrspace(13)*, i8 addrspace(13)* addrspace(11)* %117, align 16, !dbg !330, !tbaa !79, !alias.scope !81, !noalias !37, !enzyme_type !84, !enzymejl_byref_BITS_VALUE !0, !enzymejl_source_type_Ptr\7BFloat64\7D !0

Stacktrace:
  [1] setindex!
    @ ./array.jl:1021 [inlined]
  [2] fill!
    @ ./array.jl:395 [inlined]
  [3] copyto!
    @ ./broadcast.jl:964 [inlined]
  [4] materialize!
    @ ./broadcast.jl:914 [inlined]
  [5] materialize!
    @ ./broadcast.jl:911 [inlined]
  [6] eval!
    @ ./REPL[12]:4
  [7] cons
    @ ./REPL[33]:4 [inlined]
  [8] #5
    @ ./REPL[26]:3 [inlined]
  [9] fwddiffe2julia__5_6175_inner_1wrap
    @ ./REPL[26]:0
 [10] macro expansion
    @ ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:5463 [inlined]
 [11] enzyme_call
    @ ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:4997 [inlined]
 [12] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:4885 [inlined]
 [13] autodiff
    @ ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:655 [inlined]
 [14] autodiff
    @ ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:545 [inlined]
 [15] autodiff
    @ ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:517 [inlined]
 [16] (::OptimizationEnzymeExt.var"#cons_j!#41"{…})(J::Matrix{…}, θ::Vector{…})
    @ OptimizationEnzymeExt ~/.julia/packages/OptimizationBase/UXLhR/ext/OptimizationEnzymeExt.jl:232
 [17] eval_constraint_jacobian(evaluator::OptimizationMOI.MOIOptimizationNLPEvaluator{…}, j::SubArray{…}, x::Vector{…})
    @ OptimizationMOI ~/.julia/packages/OptimizationMOI/L0B28/src/nlp.jl:288
 [18] eval_constraint_jacobian(model::IpoptMathOptInterfaceExt.Optimizer, values::Vector{Float64}, x::Vector{Float64})
    @ IpoptMathOptInterfaceExt ~/.julia/packages/Ipopt/Fbuwv/ext/IpoptMathOptInterfaceExt/MOI_wrapper.jl:1248
 [19] (::IpoptMathOptInterfaceExt.var"#eval_jac_g_cb#10"{…})(x::Vector{…}, rows::Vector{…}, cols::Vector{…}, values::Vector{…})
    @ IpoptMathOptInterfaceExt ~/.julia/packages/Ipopt/Fbuwv/ext/IpoptMathOptInterfaceExt/MOI_wrapper.jl:1377
 [20] _Eval_Jac_G_CB(n::Int32, x_ptr::Ptr{…}, ::Int32, ::Int32, nele_jac::Int32, iRow::Ptr{…}, jCol::Ptr{…}, values_ptr::Ptr{…}, user_data::Ptr{…})
    @ Ipopt ~/.julia/packages/Ipopt/Fbuwv/src/C_wrapper.jl:0
 [21] (::Ipopt.var"#5#6"{IpoptProblem, Base.RefValue{Float64}})()
    @ Ipopt ~/.julia/packages/Ipopt/Fbuwv/src/C_wrapper.jl:407
 [22] disable_sigint
    @ ./c.jl:473 [inlined]
 [23] IpoptSolve
    @ ~/.julia/packages/Ipopt/Fbuwv/src/C_wrapper.jl:406 [inlined]
 [24] optimize!(model::IpoptMathOptInterfaceExt.Optimizer)
    @ IpoptMathOptInterfaceExt ~/.julia/packages/Ipopt/Fbuwv/ext/IpoptMathOptInterfaceExt/MOI_wrapper.jl:1523
 [25] __solve(cache::OptimizationMOI.MOIOptimizationNLPCache{OptimizationMOI.MOIOptimizationNLPEvaluator{…}, IpoptMathOptInterfaceExt.Optimizer})
    @ OptimizationMOI ~/.julia/packages/OptimizationMOI/L0B28/src/nlp.jl:552
 [26] solve!(cache::OptimizationMOI.MOIOptimizationNLPCache{OptimizationMOI.MOIOptimizationNLPEvaluator{…}, IpoptMathOptInterfaceExt.Optimizer})
    @ SciMLBase ~/.julia/packages/SciMLBase/iHgIu/src/solve.jl:227
 [27] solve(::OptimizationProblem{…}, ::IpoptMathOptInterfaceExt.Optimizer; kwargs::@Kwargs{})
    @ SciMLBase ~/.julia/packages/SciMLBase/iHgIu/src/solve.jl:129
 [28] solve(::OptimizationProblem{…}, ::IpoptMathOptInterfaceExt.Optimizer)
    @ SciMLBase ~/.julia/packages/SciMLBase/iHgIu/src/solve.jl:126
 [29] top-level scope
    @ REPL[38]:1
Some type information was truncated. Use `show(err)` to see complete types.

I guess this is because as @gdalle said (thanks Guillaume!), OptimizationBase is ignoring the annotations.

Also @gdalle you proposed an immediate solution through calling Enzyme through DI but in that case shouldnt I just pass my own gradients with by calling Enzyme directly? It seems like once I can not use AutoEnzyme() in the OptimizationFunction() then I might aswell pass gradients through autodiff() ?

I’m a bit confused as to what are the next steps?

Yeah in that case I’d just pass in the gradient function yourself with Enzyme.gradient or similar.

In the background I’ll make a PR adding support for the annotation properly to Optimization.jl

2 Likes

Yeah you’re right, DI won’t be much help here. Billy’s proposed fix is the right one.

There’s still something I don’t understand though: why can we use DuplicatedNoNeed to annotate a function? In the docstring it says

This should only be used if x is a write-only variable. Otherwise, if the differentiated function stores values in x and reads them back in subsequent computations, using DuplicatedNoNeed may result in incorrect derivatives. In particular, DuplicatedNoNeed should not be used for preallocated workspace, even if the user might not care about its final value, as marking a variable as NoNeed means that reads from the variable are now undefined.

Aren’t we reading from a pre allocated cache here? Why does it still work?

1 Like

I personally shy away from DuplicatedNoNeed in generic code since it creates a strong coupling between an assumption (variable is write only) and the user provided code.

As an example from my recent life experience…

If one wanted to differentiate the function below it would be appealing to use DuplicatedNoNeed on res in forward mode, and as written this is perfectly legal.

function G_Midpoint!(res, uₙ, Δt, f!, du, u, p, t; α = 0.5)
    uuₙ = (α .* uₙ .+ (1 - α) .* u)
    f!(du, uuₙ, p, t + α * Δt)

    res .= uₙ .+ Δt .* du .- u
    return nothing
end

Two weeks later a tired me wants to performance optimize that function and remove an “unnecessary” allocation, I realized that I can reuse res as a temporary variable!

function G_Midpoint!(res, uₙ, Δt, f!, du, u, p, t; α = 0.5)
    ## Use res for a temporary allocation (uₙ .+ u) ./ 2
    uuₙ = res
    uuₙ .= (α .* uₙ .+ (1 - α) .* u)
    f!(du, uuₙ, p, t + α * Δt)

    res .= uₙ .+ Δt .* du .- u
    return nothing
end

Ouf now my directional derivatives is incorrect since DuplicatedNoNeed turns off the primal calculation for uun!

Don’t want to talk about how long it took me to figure this out.

I think you’re right. My understanding was that DuplicatedNoNeed just told Enzyme that we do not care about differentiating w.r.t to this object but documentation now is very clear about its meaning and how a pre-allocated storage should not be labelled as DuplicatedNoNeed. This is better documented now that when I first came across DuplicatedNoNeed

Although I have to say that trying out both, it does not seem to matter in my application and labelling the pre-allocated storage as Duplicated or DuplicatedNoNeed produce the same gradient

1 Like

I would be really curious to understand why it seems to work in this case, and whether it is just an accident. I’ve had a similar experience where I used DuplicatedNoNeed for preallocated storage without issue, even though the documentation says we shouldn’t read again from that variable

It essentially tells Enzyme that it doesn’t need to perform the writes into the primal. This can be a non-trivial cost savings, analagous to if you specify that you don’t need the original primal return.

Whether it causes a correctness issue depends on the code. If no code reads from the buffer as we require in the docstring, it is definitionally correct. However, that doesn’t mean that reading from the buffer will cause it to be incorrect. Specifically, the issue is whether the read from the duplicatednoneed data is used in the derivative computation.

E.g. if you compute x + y, we don’t need the values of x or y. So if x or y were a load from that array , it remains correct. However, if you compute x * y, the derivative is x * dy + y * dx so then the values of x and y would be undefined and thus an incorrect derivative could arise. I also say could because if alternatively the original memory x was loaded from already had x stored there (aka you did x[1] = x[1] or similarly defined it as the same value), you’d still get the correct result.

We provide strong guarantees for correctness if you follow the spec. The behavior if you don’t follow the spec is undefined, which could still be correct – or incorrect.

2 Likes

I’m having trouble figuring out what packages you need to run your code, but the support for Duplicated / DuplicatedNoNeed was landed in OptimizationBase.jl last night so give it a go?

Thanks for the follow up. Did something change in Enzyme? I am getting an error when evaluating using Enzyme now.

julia> using Enzyme
Info Given Enzyme was explicitly requested, output will be shown live 
ERROR: LoadError: UndefVarError: `RNumber` not defined
Stacktrace:
  [1] getproperty(x::Module, f::Symbol)
    @ Base ./Base.jl:31
  [2] top-level scope
    @ ~/.julia/packages/Enzyme/hu9gq/src/analyses/activity.jl:76
  [3] include(mod::Module, _path::String)
    @ Base ./Base.jl:495
  [4] include(x::String)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:1
  [5] top-level scope
    @ ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:175
  [6] include(mod::Module, _path::String)
    @ Base ./Base.jl:495
  [7] include(x::String)
    @ Enzyme ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:1
  [8] top-level scope
    @ ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:127
  [9] include
    @ ./Base.jl:495 [inlined]
 [10] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::Nothing)
    @ Base ./loading.jl:2292
 [11] top-level scope
    @ stdin:4
in expression starting at /Users/miguelborrero/.julia/packages/Enzyme/hu9gq/src/analyses/activity.jl:76
in expression starting at /Users/miguelborrero/.julia/packages/Enzyme/hu9gq/src/compiler.jl:1
in expression starting at /Users/miguelborrero/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:1
in expression starting at stdin:4
  ✗ Enzyme
Precompiling Enzyme finished.
  0 dependencies successfully precompiled in 3 seconds. 19 already precompiled.

ERROR: The following 1 direct dependency failed to precompile:

Enzyme 

Failed to precompile Enzyme [7da242da-08ed-463a-9acd-ee780be4f1d9] to "/Users/miguelborrero/.julia/compiled/v1.10/Enzyme/jl_GxNEEA".
ERROR: LoadError: UndefVarError: `RNumber` not defined
Stacktrace:
  [1] getproperty(x::Module, f::Symbol)
    @ Base ./Base.jl:31
  [2] top-level scope
    @ ~/.julia/packages/Enzyme/hu9gq/src/analyses/activity.jl:76
  [3] include(mod::Module, _path::String)
    @ Base ./Base.jl:495
  [4] include(x::String)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:1
  [5] top-level scope
    @ ~/.julia/packages/Enzyme/hu9gq/src/compiler.jl:175
  [6] include(mod::Module, _path::String)
    @ Base ./Base.jl:495
  [7] include(x::String)
    @ Enzyme ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:1
  [8] top-level scope
    @ ~/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:127
  [9] include
    @ ./Base.jl:495 [inlined]
 [10] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::Nothing)
    @ Base ./loading.jl:2292
 [11] top-level scope
    @ stdin:4
in expression starting at /Users/miguelborrero/.julia/packages/Enzyme/hu9gq/src/analyses/activity.jl:76
in expression starting at /Users/miguelborrero/.julia/packages/Enzyme/hu9gq/src/compiler.jl:1
in expression starting at /Users/miguelborrero/.julia/packages/Enzyme/hu9gq/src/Enzyme.jl:1
in expression starting at stdin:

I’m using Julia version: Version 1.10.9 (2025-03-10) and same happens if I run it on the latest Julia version.

What’s the result of ] st

Also can you run pkg update

(Environment) pkg> st
Status `~/Documents/Optimization_julia/Environment/Project.toml`
⌃ [7da242da] Enzyme v0.13.44
⌃ [b6b21f68] Ipopt v1.10.3
  [7f7a1694] Optimization v4.3.0
  [fd9f6733] OptimizationMOI v0.5.3
  [36348300] OptimizationOptimJL v0.4.3
  [276daf66] SpecialFunctions v2.5.1
Info Packages marked with ⌃ have new versions available and may be upgradable.

After running pkg update its working. Nevertheless, this did not happen to me yesterday with v.0.13.44