Skip to content

Commit 764af82

Browse files
SebastianM-Cclaude
andcommitted
fix method ambiguities for instantiate_function in the Zygote extension
Co-authored-by: Claude <noreply@anthropic.com>
1 parent ba726da commit 764af82

File tree

1 file changed

+23
-0
lines changed

1 file changed

+23
-0
lines changed

lib/OptimizationBase/ext/OptimizationZygoteExt.jl

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,18 @@ function OptimizationBase.instantiate_function(
288288
f, x, adtype, p, num_cons; kwargs...)
289289
end
290290

291+
function OptimizationBase.instantiate_function(
292+
f::OptimizationFunction{true}, cache::OptimizationBase.ReInitCache,
293+
adtype::DifferentiationInterface.SecondOrder{
294+
<:ADTypes.AbstractADType, <:ADTypes.AutoZygote},
295+
num_cons = 0; kwargs...)
296+
x = cache.u0
297+
p = cache.p
298+
299+
return OptimizationBase.instantiate_function(
300+
f, x, adtype, p, num_cons; kwargs...)
301+
end
302+
291303
function OptimizationBase.instantiate_function(
292304
f::OptimizationFunction{true}, x,
293305
adtype::ADTypes.AutoSparse{<:Union{ADTypes.AutoZygote,
@@ -575,4 +587,15 @@ function OptimizationBase.instantiate_function(
575587
return OptimizationBase.instantiate_function(f, x, adtype, p, num_cons; kwargs...)
576588
end
577589

590+
function OptimizationBase.instantiate_function(
591+
f::OptimizationFunction{true}, cache::OptimizationBase.ReInitCache,
592+
adtype::ADTypes.AutoSparse{<:DifferentiationInterface.SecondOrder{
593+
<:ADTypes.AbstractADType, <:ADTypes.AutoZygote}},
594+
num_cons = 0; kwargs...)
595+
x = cache.u0
596+
p = cache.p
597+
598+
return OptimizationBase.instantiate_function(f, x, adtype, p, num_cons; kwargs...)
599+
end
600+
578601
end

0 commit comments

Comments
 (0)