Contexts for FiniteDiff and PolyesterForwardDiff, better context testing (#497)

* More contexts

* Stick to same contexts for same point preparation

* Doc

* Myrandom instead of zero

* Randomizable constant

Author: gdalle

DifferentiationInterface/docs/src/explanation/operators.md

@@ -130,7 +130,7 @@ For different-point preparation, the output `prep` of `prepare_op(f, b, x, [t])`
 
 For same-point preparation, the output `prep` of `prepare_op_same_point(f, b, x, [t])` can be reused in `op(f, prep, b, x, other_t)`, provided that:
 
-- the input `x` remains the same
+- the input `x` remains the same (as well as the [`Context`](@ref) constants)
 - the tangents in `t` and `other_t` have similar types and equal shapes
 
 !!! warning

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/reverse_onearg.jl

@@ -5,32 +5,49 @@ struct ChainRulesPullbackPrepSamePoint{Y,PB} <: PullbackPrep
     pb::PB
 end
 
-function DI.prepare_pullback(f, ::AutoReverseChainRules, x, ty::Tangents)
+function DI.prepare_pullback(
+    f, ::AutoReverseChainRules, x, ty::Tangents, contexts::Vararg{Constant,C}
+) where {C}
     return NoPullbackPrep()
 end
 
 function DI.prepare_pullback_same_point(
-    f, ::NoPullbackPrep, backend::AutoReverseChainRules, x, ty::Tangents
-)
+    f,
+    ::NoPullbackPrep,
+    backend::AutoReverseChainRules,
+    x,
+    ty::Tangents,
+    contexts::Vararg{Constant,C},
+) where {C}
     rc = ruleconfig(backend)
-    y, pb = rrule_via_ad(rc, f, x)
+    y, pb = rrule_via_ad(rc, f, x, map(unwrap, contexts)...)
     return ChainRulesPullbackPrepSamePoint(y, pb)
 end
 
 function DI.value_and_pullback(
-    f, ::NoPullbackPrep, backend::AutoReverseChainRules, x, ty::Tangents
-)
+    f,
+    ::NoPullbackPrep,
+    backend::AutoReverseChainRules,
+    x,
+    ty::Tangents,
+    contexts::Vararg{Constant,C},
+) where {C}
     rc = ruleconfig(backend)
-    y, pb = rrule_via_ad(rc, f, x)
+    y, pb = rrule_via_ad(rc, f, x, map(unwrap, contexts)...)
     tx = map(ty) do dy
         pb(dy)[2]
     end
     return y, tx
 end
 
 function DI.value_and_pullback(
-    f, prep::ChainRulesPullbackPrepSamePoint, ::AutoReverseChainRules, x, ty::Tangents
-)
+    f,
+    prep::ChainRulesPullbackPrepSamePoint,
+    ::AutoReverseChainRules,
+    x,
+    ty::Tangents,
+    contexts::Vararg{Constant,C},
+) where {C}
     @compat (; y, pb) = prep
     tx = map(ty) do dy
         pb(dy)[2]
@@ -39,8 +56,13 @@ function DI.value_and_pullback(
 end
 
 function DI.pullback(
-    f, prep::ChainRulesPullbackPrepSamePoint, ::AutoReverseChainRules, x, ty::Tangents
-)
+    f,
+    prep::ChainRulesPullbackPrepSamePoint,
+    ::AutoReverseChainRules,
+    x,
+    ty::Tangents,
+    contexts::Vararg{Constant,C},
+) where {C}
     @compat (; pb) = prep
     tx = map(ty) do dy
         pb(dy)[2]

DifferentiationInterface/ext/DifferentiationInterfaceFiniteDiffExt/DifferentiationInterfaceFiniteDiffExt.jl

@@ -3,6 +3,7 @@ module DifferentiationInterfaceFiniteDiffExt
 using ADTypes: AutoFiniteDiff
 import DifferentiationInterface as DI
 using DifferentiationInterface:
+    Context,
     DerivativePrep,
     GradientPrep,
     HessianPrep,
@@ -15,7 +16,9 @@ using DifferentiationInterface:
     NoJacobianPrep,
     NoPullbackPrep,
     NoPushforwardPrep,
-    Tangents
+    Tangents,
+    unwrap,
+    with_contexts
 using FiniteDiff:
     DerivativeCache,
     GradientCache,