Propagate with value types

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/differentiate_with.jl

@@ -1,7 +1,7 @@
 function ChainRulesCore.rrule(dw::DI.DifferentiateWith, x)
     (; f, backend) = dw
     y = f(x)
-    prep_same = DI.prepare_pullback_same_point(f, backend, x, (y,); strict=true)
+    prep_same = DI.prepare_pullback_same_point(f, backend, x, (y,); strict=Val(true))
     function pullbackfunc(dy)
         tx = DI.pullback(f, prep_same, backend, x, (dy,))
         return (NoTangent(), only(tx))

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/reverse_onearg.jl

@@ -1,7 +1,7 @@
 ## Pullback
 
 struct ChainRulesPullbackPrepSamePoint{SIG,Y,PB} <: DI.PullbackPrep{SIG}
-    _sig::Type{SIG}
+    _sig::Val{SIG}
     y::Y
     pb::PB
 end
@@ -12,10 +12,10 @@ function DI.prepare_pullback(
     x,
     ty::NTuple,
     contexts::Vararg{DI.GeneralizedConstant,C};
-    strict::Bool=false,
+    strict::Val=Val(false),
 ) where {C}
-    SIG = DI.signature(f, backend, x, ty, contexts...; strict)
-    return DI.NoPullbackPrep{SIG}()
+    _sig = DI.signature(f, backend, x, ty, contexts...; strict)
+    return DI.NoPullbackPrep(_sig)
 end
 
 function DI.prepare_pullback_same_point(
@@ -25,13 +25,13 @@ function DI.prepare_pullback_same_point(
     x,
     ty::NTuple,
     contexts::Vararg{DI.GeneralizedConstant,C};
-    strict::Bool=false,
+    strict::Val=Val(false),
 ) where {C}
     DI.check_prep(f, prep, backend, x, ty, contexts...)
-    SIG = DI.signature(f, backend, x, ty, contexts...; strict)
+    _sig = DI.signature(f, backend, x, ty, contexts...; strict)
     rc = ruleconfig(backend)
     y, pb = rrule_via_ad(rc, f, x, map(DI.unwrap, contexts)...)
-    return ChainRulesPullbackPrepSamePoint(SIG, y, pb)
+    return ChainRulesPullbackPrepSamePoint(_sig, y, pb)
 end
 
 function DI.value_and_pullback(

DifferentiationInterface/ext/DifferentiationInterfaceDiffractorExt/DifferentiationInterfaceDiffractorExt.jl

@@ -11,8 +11,8 @@ DI.pullback_performance(::AutoDiffractor) = DI.PullbackSlow()
 ## Pushforward
 
 function DI.prepare_pushforward(f, backend::AutoDiffractor, x, tx::NTuple)
-    SIG = DI.signature(f, backend, x, tx)
-    return DI.NoPushforwardPrep{SIG}()
+    _sig = DI.signature(f, backend, x, tx)
+    return DI.NoPushforwardPrep(_sig)
 end
 
 function DI.pushforward(

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -6,10 +6,10 @@ function DI.prepare_pushforward(
     x,
     tx::NTuple,
     contexts::Vararg{DI.Context,C};
-    strict::Bool=false,
+    strict::Val=Val(false),
 ) where {F,C}
-    SIG = DI.signature(f, backend, x, tx, contexts...; strict)
-    return DI.NoPushforwardPrep{SIG}()
+    _sig = DI.signature(f, backend, x, tx, contexts...; strict)
+    return DI.NoPushforwardPrep(_sig)
 end
 
 function DI.value_and_pushforward(
@@ -117,24 +117,22 @@ end
 ## Gradient
 
 struct EnzymeForwardGradientPrep{SIG,B,O} <: DI.GradientPrep{SIG}
+    _sig::Val{SIG}
+    _valB::Val{B}
     shadows::O
 end
 
-function EnzymeForwardGradientPrep(::Type{SIG}, ::Val{B}, shadows::O) where {SIG,B,O}
-    return EnzymeForwardGradientPrep{SIG,B,O}(shadows)
-end
-
 function DI.prepare_gradient(
     f::F,
     backend::AutoEnzyme{<:ForwardMode,<:Union{Nothing,Const}},
     x,
     contexts::Vararg{DI.Constant,C};
-    strict::Bool=false,
+    strict::Val=Val(false),
 ) where {F,C}
-    SIG = DI.signature(f, backend, x, contexts...; strict)
+    _sig = DI.signature(f, backend, x, contexts...; strict)
     valB = to_val(DI.pick_batchsize(backend, x))
     shadows = create_shadows(valB, x)
-    return EnzymeForwardGradientPrep(SIG, valB, shadows)
+    return EnzymeForwardGradientPrep(_sig, valB, shadows)
 end
 
 function DI.gradient(
@@ -199,28 +197,24 @@ end
 ## Jacobian
 
 struct EnzymeForwardOneArgJacobianPrep{SIG,B,O} <: DI.JacobianPrep{SIG}
+    _sig::Val{SIG}
+    _valB::Val{B}
     shadows::O
     output_length::Int
 end
 
-function EnzymeForwardOneArgJacobianPrep(
-    ::Type{SIG}, ::Val{B}, shadows::O, output_length::Integer
-) where {SIG,B,O}
-    return EnzymeForwardOneArgJacobianPrep{SIG,B,O}(shadows, output_length)
-end
-
 function DI.prepare_jacobian(
     f::F,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing},<:Union{Nothing,Const}},
     x,
     contexts::Vararg{DI.Constant,C};
-    strict::Bool=false,
+    strict::Val=Val(false),
 ) where {F,C}
-    SIG = DI.signature(f, backend, x, contexts...; strict)
+    _sig = DI.signature(f, backend, x, contexts...; strict)
     y = f(x, map(DI.unwrap, contexts)...)
     valB = to_val(DI.pick_batchsize(backend, x))
     shadows = create_shadows(valB, x)
-    return EnzymeForwardOneArgJacobianPrep(SIG, valB, shadows, length(y))
+    return EnzymeForwardOneArgJacobianPrep(_sig, valB, shadows, length(y))
 end
 
 function DI.jacobian(

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl

@@ -7,10 +7,10 @@ function DI.prepare_pushforward(
     x,
     tx::NTuple,
     contexts::Vararg{DI.Context,C};
-    strict::Bool=false,
+    strict::Val=Val(false),
 ) where {F,C}
-    SIG = DI.signature(f!, y, backend, x, tx, contexts...; strict)
-    return DI.NoPushforwardPrep{SIG}()
+    _sig = DI.signature(f!, y, backend, x, tx, contexts...; strict)
+    return DI.NoPushforwardPrep(_sig)
 end
 
 function DI.value_and_pushforward(

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl

@@ -48,6 +48,7 @@ end
 ## Pullback
 
 struct EnzymeReverseOneArgPullbackPrep{SIG,Y} <: DI.PullbackPrep{SIG}
+    _sig::Val{SIG}
     y_example::Y  # useful to create return activity
 end
 
@@ -57,11 +58,11 @@ function DI.prepare_pullback(
     x,
     ty::NTuple,
     contexts::Vararg{DI.Context,C};
-    strict::Bool=false,
+    strict::Val=Val(false),
 ) where {F,C}
-    SIG = DI.signature(f, backend, x, ty, contexts...; strict)
+    _sig = DI.signature(f, backend, x, ty, contexts...; strict)
     y = f(x, map(DI.unwrap, contexts)...)
-    return EnzymeReverseOneArgPullbackPrep{SIG,typeof(y)}(y)
+    return EnzymeReverseOneArgPullbackPrep(_sig, y)
 end
 
 ### Out-of-place
@@ -195,10 +196,10 @@ function DI.prepare_gradient(
     backend::AutoEnzyme{<:Union{ReverseMode,Nothing}},
     x,
     contexts::Vararg{DI.Context,C};
-    strict::Bool=false,
+    strict::Val=Val(false),
 ) where {F,C}
-    SIG = DI.signature(f, backend, x, contexts...; strict)
-    return DI.NoGradientPrep{SIG}()
+    _sig = DI.signature(f, backend, x, contexts...; strict)
+    return DI.NoGradientPrep(_sig)
 end
 
 ### Enzyme gradient API (only constants)

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_twoarg.jl

@@ -1,6 +1,7 @@
 ## Pullback
 
 struct EnzymeReverseTwoArgPullbackPrep{SIG,TY} <: DI.PullbackPrep{SIG}
+    _sig::Val{SIG}
     ty_copy::TY
 end
 
@@ -11,11 +12,11 @@ function DI.prepare_pullback(
     x,
     ty::NTuple,
     contexts::Vararg{DI.Context,C};
-    strict::Bool=false,
+    strict::Val=Val(false),
 ) where {F,C}
-    SIG = DI.signature(f!, y, backend, x, ty, contexts...; strict)
+    _sig = DI.signature(f!, y, backend, x, ty, contexts...; strict)
     ty_copy = map(copy, ty)
-    return EnzymeReverseTwoArgPullbackPrep{SIG,typeof(ty_copy)}(ty_copy)
+    return EnzymeReverseTwoArgPullbackPrep(_sig, ty_copy)
 end
 
 function DI.value_and_pullback(