Merge pull request #21 from gdalle/chainrules

Add AutoChainRules(ruleconfig)
SciML · Mar 8, 2024 · 8684021 · 8684021
2 parents f6fce4f + 2f6afad
commit 8684021
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/src/ADTypes.jl b/src/ADTypes.jl
@@ -14,6 +14,10 @@ abstract type AbstractSparseReverseMode <: AbstractReverseMode end
 abstract type AbstractSparseForwardMode <: AbstractForwardMode end
 abstract type AbstractSparseFiniteDifferences <: AbstractFiniteDifferencesMode end
 
+Base.@kwdef struct AutoChainRules{RC} <: AbstractADType
+ ruleconfig::RC
+end
+
 Base.@kwdef struct AutoFiniteDiff{T1, T2, T3} <: AbstractFiniteDifferencesMode
  fdtype::T1 = Val(:forward)
  fdjtype::T2 = fdtype
@@ -78,7 +82,8 @@ Base.@kwdef struct AutoSparseReverseDiff <: AbstractSparseReverseMode
  compile::Bool = false
 end
 
-export AutoFiniteDiff,
+export AutoChainRules,
+ AutoFiniteDiff,
  AutoFiniteDifferences,
  AutoForwardDiff,
  AutoReverseDiff,

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -4,6 +4,11 @@ using Test
 struct CustomTag end
 
 @testset "ADTypes.jl" begin
+ adtype = AutoChainRules(:ruleconfig_placeholder)
+ @test adtype isa ADTypes.AbstractADType
+ @test adtype isa AutoChainRules
+ @test adtype.ruleconfig == :ruleconfig_placeholder
+
  adtype = AutoFiniteDiff()
  @test adtype isa ADTypes.AbstractADType
  @test adtype isa AutoFiniteDiff