Add extras subtypings

JuliaDiff · Apr 25, 2024 · 26e4089 · 26e4089
1 parent 0a484ac
commit 26e4089
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 9 deletions.
diff --git a/...ationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl b/...ationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl
@@ -3,6 +3,11 @@ module DifferentiationInterfaceEnzymeExt
 using ADTypes: ADTypes, AutoEnzyme
 import DifferentiationInterface as DI
 using DifferentiationInterface:
+    DerivativeExtras,
+    GradientExtras,
+    JacobianExtras,
+    PullbackExtras,
+    PushforwardExtras,
     NoDerivativeExtras,
     NoGradientExtras,
     NoJacobianExtras,

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl
@@ -33,7 +33,7 @@ end
 
 ## Gradient
 
-struct EnzymeForwardGradientExtras{C,O}
+struct EnzymeForwardGradientExtras{C,O} <: GradientExtras
     shadow::O
 end
 
@@ -72,7 +72,7 @@ end
 
 ## Jacobian
 
-struct EnzymeForwardOneArgJacobianExtras{C,O}
+struct EnzymeForwardOneArgJacobianExtras{C,O} <: JacobianExtras
     shadow::O
 end
 

diff --git a/...erface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl b/...erface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl
@@ -3,13 +3,18 @@ module DifferentiationInterfaceFiniteDiffExt
 using ADTypes: AutoFiniteDiff
 import DifferentiationInterface as DI
 using DifferentiationInterface:
+    DerivativeExtras,
+    GradientExtras,
+    HessianExtras,
+    JacobianExtras,
+    PullbackExtras,
+    PushforwardExtras,
     NoDerivativeExtras,
     NoGradientExtras,
     NoHessianExtras,
     NoJacobianExtras,
     NoPullbackExtras,
-    NoPushforwardExtras,
-    GradientExtras
+    NoPushforwardExtras
 using FiniteDiff:
     DerivativeCache,
     GradientCache,

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/onearg.jl
@@ -19,7 +19,7 @@ end
 
 ## Derivative
 
-struct FiniteDiffOneArgDerivativeExtras{C}
+struct FiniteDiffOneArgDerivativeExtras{C} <: DerivativeExtras
     cache::C
 end
 
@@ -128,7 +128,7 @@ end
 
 ## Jacobian
 
-struct FiniteDiffOneArgJacobianExtras{C}
+struct FiniteDiffOneArgJacobianExtras{C} <: JacobianExtras
     cache::C
 end
 
@@ -166,7 +166,7 @@ end
 
 ## Hessian
 
-struct FiniteDiffHessianExtras{C}
+struct FiniteDiffHessianExtras{C} <: HessianExtras
     cache::C
 end
 

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/twoarg.jl
@@ -19,7 +19,7 @@ end
 
 ## Derivative
 
-struct FiniteDiffTwoArgDerivativeExtras{C}
+struct FiniteDiffTwoArgDerivativeExtras{C} <: DerivativeExtras
     cache::C
 end
 
@@ -61,7 +61,7 @@ end
 
 ## Jacobian
 
-struct FiniteDiffTwoArgJacobianExtras{C}
+struct FiniteDiffTwoArgJacobianExtras{C} <: JacobianExtras
     cache::C
 end