Merge pull request #10 from SciML/dw/enzyme

Support Enzyme forward + reverse mode
SciML · Jul 25, 2023 · 9e91d8a · 9e91d8a
2 parents 76daa27 + acdf895
commit 9e91d8a
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 2 deletions.
diff --git a/src/ADTypes.jl b/src/ADTypes.jl
@@ -36,7 +36,11 @@ AutoReverseDiff(; compile = false) = AutoReverseDiff(compile)
 
 struct AutoZygote <: AbstractADType end
 
-struct AutoEnzyme <: AbstractADType end
+struct AutoEnzyme{M} <: AbstractADType
+ mode::M
+end
+
+AutoEnzyme(; mode = nothing) = AutoEnzyme(mode)
 
 struct AutoTracker <: AbstractADType end
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -5,6 +5,9 @@ using Test
  adtype = AutoFiniteDiff()
  @test adtype isa ADTypes.AbstractADType
  @test adtype isa AutoFiniteDiff
+ @test adtype.fdtype === Val(:forward)
+ @test adtype.fdjtype === Val(:forward)
+ @test adtype.fdhtype === Val(:hcentral)
 
  adtype = AutoFiniteDifferences()
  @test adtype isa ADTypes.AbstractADType
@@ -19,11 +22,21 @@ using Test
 
  adtype = AutoForwardDiff()
  @test adtype isa ADTypes.AbstractADType
- @test adtype isa AutoForwardDiff
+ @test adtype isa AutoForwardDiff{nothing}
+
+ adtype = AutoForwardDiff(10)
+ @test adtype isa ADTypes.AbstractADType
+ @test adtype isa AutoForwardDiff{10}
 
  adtype = AutoReverseDiff()
  @test adtype isa ADTypes.AbstractADType
  @test adtype isa AutoReverseDiff
+ @test !adtype.compile
+
+ adtype = AutoReverseDiff(; compile = true)
+ @test adtype isa ADTypes.AbstractADType
+ @test adtype isa AutoReverseDiff
+ @test adtype.compile
 
  adtype = AutoZygote()
  @test adtype isa ADTypes.AbstractADType
@@ -32,4 +45,38 @@ using Test
  adtype = AutoTracker()
  @test adtype isa ADTypes.AbstractADType
  @test adtype isa AutoTracker
+
+ 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 AutoEnzyme{Val{:Reverse}}
+
+ adtype = AutoModelingToolkit()
+ @test adtype isa ADTypes.AbstractADType
+ @test adtype isa AutoModelingToolkit
+ @test !adtype.obj_sparse
+ @test !adtype.cons_sparse
+
+ adtype = AutoModelingToolkit(; obj_sparse = true, cons_sparse = true)
+ @test adtype isa ADTypes.AbstractADType
+ @test adtype isa AutoModelingToolkit
+ @test adtype.obj_sparse
+ @test adtype.cons_sparse
+
+ adtype = AutoSparseFiniteDiff()
+ @test adtype isa ADTypes.AbstractADType
+ @test adtype isa AutoSparseFiniteDiff
+
+ adtype = AutoSparseForwardDiff()
+ @test adtype isa ADTypes.AbstractADType
+ @test adtype isa AutoSparseForwardDiff{nothing}
+
+ adtype = AutoSparseForwardDiff(10)
+ @test adtype isa ADTypes.AbstractADType
+ @test adtype isa AutoSparseForwardDiff{10}
 end