Skip to content

Move duplicated AD helpers into main package #48

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
kshyatt wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Move duplicated AD helpers into main package #48

Show file tree
Hide file tree
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
19 changes: 1 addition & 18 deletions ext/VectorInterfaceEnzymeExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,28 +3,11 @@ module VectorInterfaceEnzymeExt
# COV_EXCL_START
# Enzyme rules aren't reachable by coverage
using VectorInterface
using VectorInterface: project_scalar, project_add!
using Enzyme
using Enzyme.EnzymeCore
using Enzyme.EnzymeCore: EnzymeRules

function project_add!(C, A, α)
TC = Base.promote_op(+, scalartype(A), scalartype(α))
return if !(TC <: Real) && scalartype(C) <: Real
add!(C, real(add!(zerovector(C, TC), A, α)))
else
add!(C, A, α)
end
end

"""
project_scalar(x::Number, dx::Number)

Project a computed tangent `dx` onto the correct tangent type for `x`.
For example, we might compute a complex `dx` but only require the real part.
"""
project_scalar(x::Number, dx::Number) = oftype(x, dx)
project_scalar(x::Real, dx::Complex) = project_scalar(x, real(dx))

function EnzymeRules.augmented_primal(
config::EnzymeRules.RevConfigWidth{1},
func::Const{typeof(scale!)},
Expand Down
21 changes: 1 addition & 20 deletions ext/VectorInterfaceMooncakeExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,31 +1,12 @@
module VectorInterfaceMooncakeExt

using VectorInterface
using VectorInterface: project_scalar, project_add!
using Mooncake
using Mooncake: @is_primitive, DefaultCtx,
NoFData, NoRData, NoTangent,
CoDual, Dual, arrayify, primal, extract

# Projection
# ----------
"""
project_scalar(x::Number, dx::Number)

Project a computed tangent `dx` onto the correct tangent type for `x`.
For example, we might compute a complex `dx` but only require the real part.
"""
project_scalar(x::Number, dx::Number) = oftype(x, dx)
project_scalar(x::Real, dx::Complex) = project_scalar(x, real(dx))

function project_add!(C, A, α)
TC = Base.promote_op(+, scalartype(A), scalartype(α))
return if !(TC <: Real) && scalartype(C) <: Real
add!(C, real(add!(zerovector(C, TC), A, α)))
else
add!(C, A, α)
end
end

_needs_tangent(x) = _needs_tangent(typeof(x))
_needs_tangent(::Type{T}) where {T <: Number} =
Mooncake.rdata_type(Mooncake.tangent_type(T)) !== NoRData
Expand Down
3 changes: 3 additions & 0 deletions src/VectorInterface.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,4 +48,7 @@ include("fallbacks.jl")
# Minimal vector type for testing
include("minimalvec.jl")

# Common AD helper methods
include("ad.jl")

end
17 changes: 17 additions & 0 deletions src/ad.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
function project_add!(C, A, α)
TC = Base.promote_op(+, scalartype(A), scalartype(α))
return if !(TC <: Real) && scalartype(C) <: Real
add!(C, real(add!(zerovector(C, TC), A, α)))

@Jutho Jutho Jun 16, 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.

The problem with this definition is of course that it is much more limited than the types that VI supports, because of the real call. Effectively, this will only support AbstractArray{<:Number}, not nested arrays like add! does.

I don't think there is a quick solution for that. One way might have to have a a project_add! implementation that is similar to add!, but that forces to also have project_add!! with implementations for tuples etc, as that is how add! is implemented.

An alternative, which might be useful in other places as well, is that VectorInterface adds a new method to its interface, namely project_real, that supports nested types. I certainly have contemplated about this in other places as well where I thought this could be useful, but then always found another way around this.

@kshyatt kshyatt Jun 16, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I suppose one option would be to require such types to implement their own project_add!, which they can at least do more easily if it's actually a "first class citizen" package method?

@Jutho Jutho Jun 16, 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.

For other types, yes, but I mean, the current implementation already fails for something like A=[(1., 2.), (3., 4.)] which is currently supported by VectorInterface.add!.

@kshyatt kshyatt Jun 16, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Then indeed maybe project_real is a good option, if nothing else so that examples like that don't blow up beneath people unexpectedly

@Jutho Jutho Jun 16, 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.

Interested to also hear @lkdvos' opinion on this.

else
add!(C, A, α)
end
end

"""
project_scalar(x::Number, dx::Number)

Project a computed tangent `dx` onto the correct tangent type for `x`.
For example, we might compute a complex `dx` but only require the real part.
"""
project_scalar(x::Number, dx::Number) = oftype(x, dx)
project_scalar(x::Real, dx::Complex) = project_scalar(x, real(dx))
Loading