Skip to content

Avoid non-public promote_op in OptimizationBase #1249

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ChrisRackauckas-Claude wants to merge 1 commit into
base: master
Choose a base branch
from
+10 −3
Draft

Avoid non-public promote_op in OptimizationBase #1249

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 10 additions & 3 deletions lib/OptimizationBase/src/OptimizationDIExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,20 +23,27 @@ using OptimizationBase.FastClosures
@inline _prep_valid(::Type{T}, v) where {T} = typeof(v) === T

# Output-buffer eltype for the `p`-accepting constraint wrapper: the type `f.cons` produces,
# including the *nested* dual when both `x` (DI's seeds) and `p` (the sensitivity layer) carry
# duals with different tags (`promote_op(+, …)` nests them). Deliberately uses plain `eltype(p)`
# including nested duals when both `x` (DI's seeds) and `p` (the sensitivity layer) carry
# duals. Deliberately uses plain `eltype(p)`
# rather than `SciMLBase.anyeltypedual`: this runs *inside* the DI-differentiated wrapper, and
# Enzyme's forward mode corrupts the derivative shadow (DataType-valued entries) when the
# allocation type flows through `anyeltypedual` — in either its value or its type-level form,
# and even with an `EnzymeRules.inactive` mark on this helper. A structured `p` (non-`Number`
# eltype) therefore falls back to `eltype(x)`; duals nested inside such a `p` are the one
# unsupported case. (`_prep_valid` above runs in the outer closure, outside anything a backend
# differentiates, so its type comparison is safe there.)
@inline function _add_output_eltype(::Type{Tu}, ::Type{Tp}) where {Tu, Tp}
return try
typeof(zero(Tu) + zero(Tp))
catch
Any
end
end
@inline function _cons_out_eltype(x, p)
Tu = eltype(x)
p isa Union{SciMLBase.NullParameters, Nothing} && return Tu
Tp = eltype(p)
return Tp <: Number ? Base.promote_op(+, Tu, Tp) : Tu
return Tp <: Number ? _add_output_eltype(Tu, Tp) : Tu
end

function instantiate_function(
Expand Down
Loading