Tapir mutation, take 2 (#141)

* Tapir mutation, take 2

* Handle the case where `f!` is a closure

* Define zero_tangent

* Copy dy

* Reactivate tests

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/DifferentiationInterfaceTapirExt.jl

@@ -4,43 +4,19 @@ using ADTypes: ADTypes
 import DifferentiationInterface as DI
 using DifferentiationInterface: AutoTapir
 using DifferentiationInterface: PullbackExtras
-using Tapir: CoDual, build_rrule, value_and_pullback!!, zero_codual
-
-DI.supports_mutation(::AutoTapir) = DI.MutationNotSupported()
-
-function zero_sametype!!(x_target, x::Number)
-    return zero(x)
-end
-
-function zero_sametype!!(x_target, x::AbstractArray)
-    x_sametype = convert(typeof(x), x_target)
-    x_sametype .= zero(eltype(x))
-    return x_sametype
-end
-
-## Pullback
-
-struct TapirPullbackExtras{R} <: PullbackExtras
-    rrule::R
-end
-
-DI.prepare_pullback(f, ::AutoTapir, x) = TapirPullbackExtras(build_rrule(f, x))
-
-function DI.value_and_pullback(f, ::AutoTapir, x, dy, extras::TapirPullbackExtras)
-    y = f(x)
-    dy_righttype = convert(typeof(y), dy)
-    _, (_, dx) = value_and_pullback!!(extras.rrule, dy_righttype, f, x)
-    return y, dx
-end
-
-function DI.value_and_pullback!!(f, dx, ::AutoTapir, x, dy, extras::TapirPullbackExtras)
-    y = f(x)
-    dy_righttype = convert(typeof(y), dy)
-    dx_righttype = zero_sametype!!(dx, x)
-    new_y, (_, new_dx) = value_and_pullback!!(
-        extras.rrule, dy_righttype, zero_codual(f), CoDual(x, dx_righttype)
-    )
-    return new_y, new_dx
-end
+using Tapir:
+    CoDual,
+    NoTangent,
+    build_rrule,
+    increment!!,
+    set_to_zero!!,
+    tangent,
+    tangent_type,
+    value_and_pullback!!,
+    zero_codual,
+    zero_tangent
+
+include("allocating.jl")
+include("mutating.jl")
 
 end

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/allocating.jl

@@ -0,0 +1,25 @@
+struct TapirAllocatingPullbackExtras{R} <: PullbackExtras
+    rrule::R
+end
+
+DI.prepare_pullback(f, ::AutoTapir, x) = TapirAllocatingPullbackExtras(build_rrule(f, x))
+
+function DI.value_and_pullback(f, ::AutoTapir, x, dy, extras::TapirAllocatingPullbackExtras)
+    y = f(x)  # TODO: one call too many, just for the conversion
+    dy_righttype = convert(tangent_type(typeof(y)), dy)
+    new_y, (new_df, new_dx) = value_and_pullback!!(extras.rrule, dy_righttype, f, x)
+    return new_y, new_dx
+end
+
+function DI.value_and_pullback!!(
+    f, dx, ::AutoTapir, x, dy, extras::TapirAllocatingPullbackExtras
+)
+    y = f(x)  # TODO: one call too many, just for the conversion
+    dy_righttype = convert(tangent_type(typeof(y)), dy)
+    dx_righttype = convert(tangent_type(typeof(x)), dx)
+    dx_righttype = set_to_zero!!(dx_righttype)
+    new_y, (new_df, new_dx) = value_and_pullback!!(
+        extras.rrule, dy_righttype, zero_codual(f), CoDual(x, dx_righttype)
+    )
+    return new_y, new_dx
+end

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/mutating.jl

@@ -0,0 +1,50 @@
+struct TapirMutatingPullbackExtras{R} <: PullbackExtras
+    rrule::R
+end
+
+function DI.prepare_pullback(f!, ::AutoTapir, y, x)
+    return TapirMutatingPullbackExtras(build_rrule(f!, y, x))
+end
+
+# see https://github.com/withbayes/Tapir.jl/issues/113#issuecomment-2036718992
+
+function DI.value_and_pullback!!(
+    f!, y, dx, ::AutoTapir, x, dy, extras::TapirMutatingPullbackExtras
+)
+    dy_righttype = convert(tangent_type(typeof(y)), copy(dy))
+    dx_righttype = convert(tangent_type(typeof(x)), dx)
+
+    # We want the VJP, not VJP + dx, so I'm going to zero-out `dx`. `set_to_zero!!` has the advantage
+    # that it will also replace any immutable components of `dx` to zero.
+    dx_righttype = set_to_zero!!(dx_righttype)
+
+    # We want `dy` to correspond to the cotangent of `y` _after_
+    # running the forwards-pass, so I'm going to take a copy, and zero-out the original.
+    dy_righttype_backup = copy(dy_righttype)
+    dy_righttype = set_to_zero!!(dy_righttype)
+
+    # Mutate a copy of `y`, so that we can run the reverse-pass later on.
+    y_copy = copy(y)
+
+    # In case `f!` is a closure
+    df! = zero_tangent(f!)
+
+    # Run the forwards-pass.
+    out, pb!! = extras.rrule(
+        CoDual(f!, df!), CoDual(y_copy, dy_righttype), CoDual(x, dx_righttype)
+    )
+
+    # Verify that the output is non-differentiable.
+    @assert tangent(out) == NoTangent()
+
+    # Set the cotangent of `y` to be equal to the requested value.
+    dy_righttype = increment!!(dy_righttype, dy_righttype_backup)
+
+    # Record the state of `y` before running the reverse-pass.
+    y = copy!(y, y_copy)
+
+    # Run the reverse-pass.
+    _, _, new_dx = pb!!(NoTangent(), df!, dy_righttype, dx_righttype)
+
+    return y, new_dx
+end