Replace Tangents with NTuple (#501)

* Replace Tangents with NTuple

* Improve Mooncake

* Second der type stability

* Values for Enzyme

* Fixes

* Fix

* Fix type params

* Fix

Author: gdalle

DifferentiationInterface/docs/src/api.md

@@ -13,7 +13,6 @@ DifferentiationInterface
 ```@docs
 Context
 Constant
-Tangents
 ```
 
 ## First order

DifferentiationInterface/docs/src/explanation/operators.md

@@ -33,7 +33,7 @@ These operators are computed using only the input `x`.
 
 ### Low-level operators
 
-These operators are computed using the input `x` and another argument `t` of type [`Tangents`](@ref), which contains one or more tangents.
+These operators are computed using the input `x` and another argument `t` of type `NTuple`, which contains one or more tangents.
 You can think of tangents as perturbations propagated through the function; they live either in the same space as `x` or in the same space as `y`.
 
 | operator                    | order | input `x` | output `y` | element type of `t` | operator result type | operator result shape |

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl

@@ -12,13 +12,7 @@ using ChainRulesCore:
 using Compat
 import DifferentiationInterface as DI
 using DifferentiationInterface:
-    Constant,
-    DifferentiateWith,
-    NoPullbackPrep,
-    NoPushforwardPrep,
-    PullbackPrep,
-    Tangents,
-    unwrap
+    Constant, DifferentiateWith, NoPullbackPrep, NoPushforwardPrep, PullbackPrep, unwrap
 
 ruleconfig(backend::AutoChainRules) = backend.ruleconfig
 

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/differentiate_with.jl

@@ -1,9 +1,9 @@
 function ChainRulesCore.rrule(dw::DifferentiateWith, x)
     @compat (; f, backend) = dw
     y = f(x)
-    prep_same = DI.prepare_pullback_same_point(f, backend, x, DI.Tangents(y))
+    prep_same = DI.prepare_pullback_same_point(f, backend, x, (y,))
     function pullbackfunc(dy)
-        tx = DI.pullback(f, prep_same, backend, x, DI.Tangents(dy))
+        tx = DI.pullback(f, prep_same, backend, x, (dy,))
         return (NoTangent(), only(tx))
     end
     return y, pullbackfunc

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/reverse_onearg.jl

@@ -6,7 +6,7 @@ struct ChainRulesPullbackPrepSamePoint{Y,PB} <: PullbackPrep
 end
 
 function DI.prepare_pullback(
-    f, ::AutoReverseChainRules, x, ty::Tangents, contexts::Vararg{Constant,C}
+    f, ::AutoReverseChainRules, x, ty::NTuple, contexts::Vararg{Constant,C}
 ) where {C}
     return NoPullbackPrep()
 end
@@ -16,7 +16,7 @@ function DI.prepare_pullback_same_point(
     ::NoPullbackPrep,
     backend::AutoReverseChainRules,
     x,
-    ty::Tangents,
+    ty::NTuple,
     contexts::Vararg{Constant,C},
 ) where {C}
     rc = ruleconfig(backend)
@@ -29,7 +29,7 @@ function DI.value_and_pullback(
     ::NoPullbackPrep,
     backend::AutoReverseChainRules,
     x,
-    ty::Tangents,
+    ty::NTuple,
     contexts::Vararg{Constant,C},
 ) where {C}
     rc = ruleconfig(backend)
@@ -45,7 +45,7 @@ function DI.value_and_pullback(
     prep::ChainRulesPullbackPrepSamePoint,
     ::AutoReverseChainRules,
     x,
-    ty::Tangents,
+    ty::NTuple,
     contexts::Vararg{Constant,C},
 ) where {C}
     @compat (; y, pb) = prep
@@ -60,7 +60,7 @@ function DI.pullback(
     prep::ChainRulesPullbackPrepSamePoint,
     ::AutoReverseChainRules,
     x,
-    ty::Tangents,
+    ty::NTuple,
     contexts::Vararg{Constant,C},
 ) where {C}
     @compat (; pb) = prep

DifferentiationInterface/ext/DifferentiationInterfaceDiffractorExt/DifferentiationInterfaceDiffractorExt.jl

@@ -2,7 +2,7 @@ module DifferentiationInterfaceDiffractorExt
 
 using ADTypes: ADTypes, AutoDiffractor
 import DifferentiationInterface as DI
-using DifferentiationInterface: NoPushforwardPrep, Tangents
+using DifferentiationInterface: NoPushforwardPrep
 using Diffractor: DiffractorRuleConfig, TaylorTangentIndex, ZeroBundle, bundle, ∂☆
 
 DI.check_available(::AutoDiffractor) = true
@@ -11,19 +11,20 @@ DI.pullback_performance(::AutoDiffractor) = DI.PullbackSlow()
 
 ## Pushforward
 
-DI.prepare_pushforward(f, ::AutoDiffractor, x, tx::Tangents) = NoPushforwardPrep()
+DI.prepare_pushforward(f, ::AutoDiffractor, x, tx::NTuple) = NoPushforwardPrep()
 
-function DI.pushforward(f, ::NoPushforwardPrep, ::AutoDiffractor, x, tx::Tangents)
+function DI.pushforward(f, ::NoPushforwardPrep, ::AutoDiffractor, x, tx::NTuple)
     ty = map(tx) do dx
         # code copied from Diffractor.jl
         z = ∂☆{1}()(ZeroBundle{1}(f), bundle(x, dx))
         dy = z[TaylorTangentIndex(1)]
+        dy
     end
     return ty
 end
 
 function DI.value_and_pushforward(
-    f, prep::NoPushforwardPrep, backend::AutoDiffractor, x, tx::Tangents
+    f, prep::NoPushforwardPrep, backend::AutoDiffractor, x, tx::NTuple
 )
     return f(x), DI.pushforward(f, prep, backend, x, tx)
 end

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -17,7 +17,6 @@ using DifferentiationInterface:
     NoJacobianPrep,
     NoPullbackPrep,
     NoPushforwardPrep,
-    Tangents,
     pick_batchsize
 using Enzyme:
     Active,

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -4,7 +4,7 @@ function DI.prepare_pushforward(
     f::F,
     ::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents,
+    tx::NTuple,
     contexts::Vararg{Context,C},
 ) where {F,C}
     return NoPushforwardPrep()
@@ -15,7 +15,7 @@ function DI.value_and_pushforward(
     ::NoPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents{1},
+    tx::NTuple{1},
     contexts::Vararg{Context,C},
 ) where {F,C}
     f_and_df = get_f_and_df(f, backend)
@@ -24,32 +24,32 @@ function DI.value_and_pushforward(
     dy, y = autodiff(
         forward_mode_withprimal(backend), f_and_df, x_and_dx, map(translate, contexts)...
     )
-    return y, Tangents(dy)
+    return y, (dy,)
 end
 
 function DI.value_and_pushforward(
     f::F,
     ::NoPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents{B},
+    tx::NTuple{B},
     contexts::Vararg{Context,C},
 ) where {F,B,C}
     f_and_df = get_f_and_df(f, backend, Val(B))
-    dxs_sametype = map(Fix1(convert, typeof(x)), tx.d)
-    x_and_dxs = BatchDuplicated(x, dxs_sametype)
-    dys, y = autodiff(
-        forward_mode_withprimal(backend), f_and_df, x_and_dxs, map(translate, contexts)...
+    tx_sametype = map(Fix1(convert, typeof(x)), tx)
+    x_and_tx = BatchDuplicated(x, tx_sametype)
+    ty, y = autodiff(
+        forward_mode_withprimal(backend), f_and_df, x_and_tx, map(translate, contexts)...
     )
-    return y, Tangents(dys...)
+    return y, values(ty)
 end
 
 function DI.pushforward(
     f::F,
     ::NoPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents{1},
+    tx::NTuple{1},
     contexts::Vararg{Context,C},
 ) where {F,C}
     f_and_df = get_f_and_df(f, backend)
@@ -60,53 +60,56 @@ function DI.pushforward(
             forward_mode_noprimal(backend), f_and_df, x_and_dx, map(translate, contexts)...
         ),
     )
-    return Tangents(dy)
+    return (dy,)
 end
 
 function DI.pushforward(
     f::F,
     ::NoPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents{B},
+    tx::NTuple{B},
     contexts::Vararg{Context,C},
 ) where {F,B,C}
     f_and_df = get_f_and_df(f, backend, Val(B))
-    dxs_sametype = map(Fix1(convert, typeof(x)), tx.d)
-    x_and_dxs = BatchDuplicated(x, dxs_sametype)
-    dys = only(
+    tx_sametype = map(Fix1(convert, typeof(x)), tx)
+    x_and_tx = BatchDuplicated(x, tx_sametype)
+    ty = only(
         autodiff(
-            forward_mode_noprimal(backend), f_and_df, x_and_dxs, map(translate, contexts)...
+            forward_mode_noprimal(backend), f_and_df, x_and_tx, map(translate, contexts)...
         ),
     )
-    return Tangents(dys...)
+    return values(ty)
 end
 
 function DI.value_and_pushforward!(
     f::F,
-    ty::Tangents,
+    ty::NTuple,
     prep::NoPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents,
+    tx::NTuple,
     contexts::Vararg{Context,C},
 ) where {F,C}
     # dy cannot be passed anyway
     y, new_ty = DI.value_and_pushforward(f, prep, backend, x, tx, contexts...)
-    return y, copyto!(ty, new_ty)
+    foreach(copyto!, ty, new_ty)
+    return y, ty
 end
 
 function DI.pushforward!(
     f::F,
-    ty::Tangents,
+    ty::NTuple,
     prep::NoPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents,
+    tx::NTuple,
     contexts::Vararg{Context,C},
 ) where {F,C}
     # dy cannot be passed anyway
-    return copyto!(ty, DI.pushforward(f, prep, backend, x, tx, contexts...))
+    new_ty = DI.pushforward(f, prep, backend, x, tx, contexts...)
+    foreach(copyto!, ty, new_ty)
+    return ty
 end
 
 ## Gradient

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl

@@ -5,7 +5,7 @@ function DI.prepare_pushforward(
     y,
     ::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents,
+    tx::NTuple,
     contexts::Vararg{Context,C},
 ) where {F,C}
     return NoPushforwardPrep()
@@ -17,7 +17,7 @@ function DI.value_and_pushforward(
     ::NoPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents{1},
+    tx::NTuple{1},
     contexts::Vararg{Context,C},
 ) where {F,C}
     f!_and_df! = get_f_and_df(f!, backend)
@@ -33,7 +33,7 @@ function DI.value_and_pushforward(
         x_and_dx,
         map(translate, contexts)...,
     )
-    return y, Tangents(dy_sametype)
+    return y, (dy_sametype,)
 end
 
 function DI.value_and_pushforward(
@@ -42,21 +42,64 @@ function DI.value_and_pushforward(
     ::NoPushforwardPrep,
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
     x,
-    tx::Tangents{B},
+    tx::NTuple{B},
     contexts::Vararg{Context,C},
 ) where {F,B,C}
     f!_and_df! = get_f_and_df(f!, backend, Val(B))
-    dxs_sametype = map(Fix1(convert, typeof(x)), tx.d)
-    dys_sametype = ntuple(_ -> make_zero(y), Val(B))
-    x_and_dxs = BatchDuplicated(x, dxs_sametype)
-    y_and_dys = BatchDuplicated(y, dys_sametype)
+    tx_sametype = map(Fix1(convert, typeof(x)), tx)
+    ty_sametype = ntuple(_ -> make_zero(y), Val(B))
+    x_and_tx = BatchDuplicated(x, tx_sametype)
+    y_and_ty = BatchDuplicated(y, ty_sametype)
     autodiff(
         forward_mode_noprimal(backend),
         f!_and_df!,
         Const,
-        y_and_dys,
-        x_and_dxs,
+        y_and_ty,
+        x_and_tx,
         map(translate, contexts)...,
     )
-    return y, Tangents(dys_sametype...)
+    return y, ty_sametype
+end
+
+function DI.pushforward(
+    f!::F,
+    y,
+    prep::NoPushforwardPrep,
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    x,
+    tx::NTuple,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    _, ty = DI.value_and_pushforward(f!, y, prep, backend, x, tx, contexts...)
+    return ty
+end
+
+function DI.value_and_pushforward!(
+    f!::F,
+    y,
+    ty::NTuple,
+    prep::NoPushforwardPrep,
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    x,
+    tx::NTuple,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    y, new_ty = DI.value_and_pushforward(f!, y, prep, backend, x, tx, contexts...)
+    foreach(copyto!, ty, new_ty)
+    return y, ty
+end
+
+function DI.pushforward!(
+    f!::F,
+    y,
+    ty::NTuple,
+    prep::NoPushforwardPrep,
+    backend::AutoEnzyme{<:Union{ForwardMode,Nothing}},
+    x,
+    tx::NTuple,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    new_ty = DI.pushforward(f!, y, prep, backend, x, tx, contexts...)
+    foreach(copyto!, ty, new_ty)
+    return ty
 end