Skip to content

Support non-Array AbstractVectors for AutoEnzyme hess and fgh! #1274

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

Open
wsmoses wants to merge 7 commits into
base: master
Choose a base branch
from
+25 −3
Open

Support non-Array AbstractVectors for AutoEnzyme hess and fgh! #1274

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e067321
Support non-Array AbstractVectors for AutoEnzyme hess and fgh!
wsmoses Jul 19, 2026
7db3a5f
Copy G_arr back into G when G is not an Array in fgh!
wsmoses Jul 19, 2026
9b59807
Use G isa Array check and copyto! in fgh!
wsmoses Jul 19, 2026
cc12a1b
Remove make_zero! in fgh! and keep copyto! for non-Array G
wsmoses Jul 19, 2026
27c8188
Restore Enzyme.make_zero!(G_arr) in fgh!
wsmoses Jul 19, 2026
bccf37d
Guard make_zero! with if !(G isa Array) in fgh!
wsmoses Jul 19, 2026
6a2fb0d
Always call Enzyme.make_zero!(G_arr) in fgh! to fix CI test failure
wsmoses Jul 20, 2026
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
14 changes: 11 additions & 3 deletions lib/OptimizationBase/ext/OptimizationEnzymeExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,12 +167,13 @@ function OptimizationBase.instantiate_function(
function hess(res, θ, p = p)
Enzyme.make_zero!(bθ)
Enzyme.make_zero!.(vdbθ)
θ_arr = θ isa Array ? θ : Array(θ)

@ChrisRackauckas ChrisRackauckas Jul 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This isn't correct, it's coercing into an array which for example would not be the correct behavior on GPUArrays.

@wsmoses wsmoses Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is the existing behavior here:

        vdθ = Tuple((Array(r) for r in eachrow(I(length(x)) * one(eltype(x)))))

already does the conversion [we just need to match this for both primal+shadow]

@ChrisRackauckas ChrisRackauckas Jul 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

well then that's the incorrect part that should get fixed.

@wsmoses wsmoses Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how? you cannot guarantee the primal/shadow are the same type [and index offsets] in a nice way that I can think of offhand otherwise (ofc longer term you can make a helper to adapt the array type, but that's a follow up)

@ChrisRackauckas ChrisRackauckas Jul 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

make_zero, Adapt.jl, ArrayInterface.restructure

@wsmoses wsmoses Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think any of those work here, but I'll let you continue it


Enzyme.autodiff(
fmode,
inner_grad,
Const(rmode),
Enzyme.BatchDuplicated(θ, vdθ),
Enzyme.BatchDuplicated(θ_arr, vdθ),
Enzyme.BatchDuplicatedNoNeed(bθ, vdbθ),
Const(f.f),
Const(p)
Expand All @@ -193,17 +194,24 @@ function OptimizationBase.instantiate_function(
function fgh!(G, H, θ, p = p)
vdθ = Tuple((Array(r) for r in eachrow(I(length(θ)) * one(eltype(θ)))))
vdbθ = Tuple(zeros(eltype(θ), length(θ)) for i in eachindex(θ))
θ_arr = θ isa Array ? θ : Array(θ)
G_arr = G isa Array ? G : Array(G)
Enzyme.make_zero!(G_arr)

Enzyme.autodiff(
fmode,
inner_grad,
Const(rmode),
Enzyme.BatchDuplicated(θ, vdθ),
Enzyme.BatchDuplicatedNoNeed(G, vdbθ),
Enzyme.BatchDuplicated(θ_arr, vdθ),
Enzyme.BatchDuplicatedNoNeed(G_arr, vdbθ),
Const(f.f),
Const(p)
)

if !(G isa Array)
copyto!(G, G_arr)
end

for i in eachindex(θ)
H[i, :] .= vdbθ[i]
end
Expand Down
14 changes: 14 additions & 0 deletions lib/OptimizationBase/test/AD/adtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,20 @@ optprob.cons_h(H3, x0)
optprob.lag_h(H4, x0, σ, μ)
@test H4 ≈ σ * H2 + μ[1] * H3[1] rtol = 1.0e-6

# Test non-Vector AbstractVector (e.g. SubArray) for AutoEnzyme hess and fgh!
x_view = @view zeros(4)[1:2]
optprob_view = OptimizationBase.instantiate_function(
OptimizationFunction(rosenbrock, OptimizationBase.AutoEnzyme()), x_view,
OptimizationBase.AutoEnzyme(), nothing, 0, h = true, fgh = true
)
H_view = Array{Float64}(undef, 2, 2)
G_view = Array{Float64}(undef, 2)
optprob_view.hess(H_view, x_view)
@test H1 == H_view
optprob_view.fgh(G_view, H_view, x_view)
@test G1 == G_view
@test H1 == H_view

@ChrisRackauckas ChrisRackauckas Jul 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should test other things like GPUArrays, StaticArrays, ComponentArrays, etc.

@wsmoses wsmoses Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'll let you take over this PR then

@ChrisRackauckas ChrisRackauckas Jul 20, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that's fine.

G2 = Array{Float64}(undef, 2)
H2 = Array{Float64}(undef, 2, 2)

Expand Down
Loading