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

Add support for custom tags to ForwardDiff types #12

Merged
merged 5 commits into from
Jul 25, 2023
Merged

Add support for custom tags to ForwardDiff types #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
93f3824
Add support for custom tags to ForwardDiff types
devmotion Jul 19, 2023
13c2bc8
Merge branch 'main' into dw/forwarddiff
Vaibhavdixit02 Jul 25, 2023
af11539
Update Project.toml
Vaibhavdixit02 Jul 25, 2023
29b7c8a
Keep nonkwarg constructor for nonbreakability
Vaibhavdixit02 Jul 25, 2023
b3ae9ac
Drop nonkwarg arg
Vaibhavdixit02 Jul 25, 2023
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
16 changes: 10 additions & 6 deletions src/ADTypes.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,12 @@ end

AutoFiniteDifferences(; fdm = nothing) = AutoFiniteDifferences(fdm)

struct AutoForwardDiff{chunksize} <: AbstractADType end
struct AutoForwardDiff{chunksize,T} <: AbstractADType
tag::T
end

function AutoForwardDiff(chunksize = nothing)
AutoForwardDiff{chunksize}()
function AutoForwardDiff(; chunksize = nothing, tag = nothing)
AutoForwardDiff{chunksize,typeof(tag)}(tag)
end

struct AutoReverseDiff <: AbstractADType
Expand Down Expand Up @@ -55,10 +57,12 @@ end

struct AutoSparseFiniteDiff <: AbstractADType end

struct AutoSparseForwardDiff{chunksize} <: AbstractADType end
struct AutoSparseForwardDiff{chunksize,T} <: AbstractADType
tag::T
end

function AutoSparseForwardDiff(chunksize = nothing)
AutoSparseForwardDiff{chunksize}()
function AutoSparseForwardDiff(; chunksize = nothing, tag = nothing)
AutoSparseForwardDiff{chunksize,typeof(tag)}(tag)
end

export AutoFiniteDiff, AutoFiniteDifferences, AutoForwardDiff, AutoReverseDiff, AutoZygote, AutoEnzyme, AutoTracker, AutoModelingToolkit, AutoSparseFiniteDiff, AutoSparseForwardDiff
Expand Down
29 changes: 18 additions & 11 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
using ADTypes
using Test

struct CustomTag end

@testset "ADTypes.jl" begin
adtype = AutoFiniteDiff()
@test adtype isa ADTypes.AbstractADType
Expand All @@ -22,11 +24,12 @@ using Test

adtype = AutoForwardDiff()
@test adtype isa ADTypes.AbstractADType
@test adtype isa AutoForwardDiff{nothing}

adtype = AutoForwardDiff(10)
@test adtype isa AutoForwardDiff{nothing,Nothing}

adtype = AutoForwardDiff(; chunksize = 10, tag = CustomTag())
@test adtype isa ADTypes.AbstractADType
@test adtype isa AutoForwardDiff{10}
@test adtype isa AutoForwardDiff{10,CustomTag}

adtype = AutoReverseDiff()
@test adtype isa ADTypes.AbstractADType
Expand All @@ -46,15 +49,13 @@ using Test
@test adtype isa ADTypes.AbstractADType
@test adtype isa AutoTracker

adtype = AutoEnzyme()
adtype = AutoSparseForwardDiff()
@test adtype isa ADTypes.AbstractADType
@test adtype isa AutoEnzyme{Nothing}
@test adtype isa AutoSparseForwardDiff{nothing,Nothing}

# In practice, you would rather specify a
# `mode::Enzyme.Mode`, e.g. `Enzyme.Reverse` or `Enzyme.Forward`
adtype = AutoEnzyme(; mode = Val(:Reverse))
adtype = AutoSparseForwardDiff(; chunksize = 10, tag = CustomTag())
@test adtype isa ADTypes.AbstractADType
@test adtype isa AutoEnzyme{Val{:Reverse}}
@test adtype isa AutoSparseForwardDiff{10,CustomTag}

adtype = AutoModelingToolkit()
@test adtype isa ADTypes.AbstractADType
Expand All @@ -76,7 +77,13 @@ using Test
@test adtype isa ADTypes.AbstractADType
@test adtype isa AutoSparseForwardDiff{nothing}

adtype = AutoSparseForwardDiff(10)
adtype = AutoEnzyme()
@test adtype isa ADTypes.AbstractADType
@test adtype isa AutoEnzyme{Nothing}

# In practice, you would rather specify a
# `mode::Enzyme.Mode`, e.g. `Enzyme.Reverse` or `Enzyme.Forward`
adtype = AutoEnzyme(; mode = Val(:Reverse))
@test adtype isa ADTypes.AbstractADType
@test adtype isa AutoSparseForwardDiff{10}
@test adtype isa AutoEnzyme{Val{:Reverse}}
end