Skip to content
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

Jump to bottom

Tag depends on ForwardDiff backend object #229

Merged
merged 2 commits into from
Apr 29, 2024
Merged

Tag depends on ForwardDiff backend object #229

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9f7157e
Tag depends on ForwardDiff backend object
gdalle Apr 29, 2024
ca92d31
Remove map in make_dual
gdalle Apr 29, 2024
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
1 change: 1 addition & 0 deletions ...face/ext/DifferentiationInterfaceForwardDiffExt/DifferentiationInterfaceForwardDiffExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
module DifferentiationInterfaceForwardDiffExt

using ADTypes: AbstractADType, AutoForwardDiff
using Base: Fix1
import DifferentiationInterface as DI
using DifferentiationInterface:
DerivativeExtras,
Expand Down
4 changes: 2 additions & 2 deletions DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,8 @@ struct ForwardDiffOneArgPushforwardExtras{T,X} <: PushforwardExtras
xdual_tmp::X
end

function DI.prepare_pushforward(f, ::AutoForwardDiff, x, dx)
T = tag_type(f, x)
function DI.prepare_pushforward(f, backend::AutoForwardDiff, x, dx)
T = tag_type(f, backend, x)
xdual_tmp = make_dual(T, x, dx)
return ForwardDiffOneArgPushforwardExtras{T,typeof(xdual_tmp)}(xdual_tmp)
end
Expand Down
4 changes: 2 additions & 2 deletions DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ struct ForwardDiffTwoArgPushforwardExtras{T,X,Y} <: PushforwardExtras
ydual_tmp::Y
end

function DI.prepare_pushforward(f!, y, ::AutoForwardDiff, x, dx)
T = tag_type(f!, x)
function DI.prepare_pushforward(f!, y, backend::AutoForwardDiff, x, dx)
T = tag_type(f!, backend, x)
xdual_tmp = make_dual(T, x, dx)
ydual_tmp = make_dual(T, y, similar(y))
return ForwardDiffTwoArgPushforwardExtras{T,typeof(xdual_tmp),typeof(ydual_tmp)}(
Expand Down
37 changes: 9 additions & 28 deletions DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,41 +1,22 @@
choose_chunk(::AutoForwardDiff{nothing}, x) = Chunk(x)
choose_chunk(::AutoForwardDiff{C}, x) where {C} = Chunk{C}()

tag_type(::F, x::Number) where {F} = Tag{F,typeof(x)}
tag_type(::F, x::AbstractArray) where {F} = Tag{F,eltype(x)}
tag_type(f, ::AutoForwardDiff{C,T}, x) where {C,T} = T
tag_type(f, ::AutoForwardDiff{C,Nothing}, x) where {C} = Tag{typeof(f),eltype(x)}

make_dual(::Type{T}, x::Number, dx) where {T} = Dual{T}(x, dx)
make_dual(::Type{T}, x::AbstractArray, dx) where {T} = Dual{T}.(x, dx)
make_dual(::Type{T}, x, dx) where {T} = Dual{T}.(x, dx) # TODO: map causes Enzyme to fail

function make_dual!(::Type{T}, xdual, x::AbstractArray, dx) where {T}
for i in eachindex(xdual, x, dx)
xdual[i] = Dual{T}(x[i], dx[i])
end
return nothing
end
make_dual!(::Type{T}, xdual, x, dx) where {T} = map!(Dual{T}, xdual, x, dx)

myvalue(::Type{T}, ydual::Number) where {T} = value(T, ydual)
myvalue(::Type{T}, ydual::AbstractArray) where {T} = value.(T, ydual)
myvalue(::Type{T}, ydual) where {T} = map(Fix1(value, T), ydual)

function myvalue!(::Type{T}, y::AbstractArray, ydual) where {T}
for i in eachindex(y, ydual)
y[i] = value(T, ydual[i])
end
return nothing
end
myvalue!(::Type{T}, y, ydual) where {T} = map!(Fix1(value, T), y, ydual)

myderivative(::Type{T}, ydual::Number) where {T} = extract_derivative(T, ydual)
myderivative(::Type{T}, ydual::AbstractArray) where {T} = extract_derivative(T, ydual)

function myderivative!(::Type{T}, dy, ydual::AbstractArray) where {T}
extract_derivative!(T, dy, ydual)
return nothing
end
myderivative(::Type{T}, ydual) where {T} = map(Fix1(extract_derivative, T), ydual)

function myvalueandderivative!(::Type{T}, y, dy, ydual::AbstractArray) where {T}
for i in eachindex(y, dy, ydual)
y[i] = value(T, ydual[i])
dy[i] = extract_derivative(T, ydual[i])
end
return nothing
function myderivative!(::Type{T}, dy, ydual) where {T}
return map!(Fix1(extract_derivative, T), dy, ydual)
end
Loading