Updates for Tapir (#241)

* Bump Tapir version

* Update Tapir usage

* Bump patch

* Update DifferentiationInterface/Project.toml

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

---------

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

Author: willtebbutt

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -57,7 +57,7 @@ PolyesterForwardDiff = "0.1.1"
 ReverseDiff = "1.15.1"
 SparseArrays = "1"
 Symbolics = "5.27.1"
-Tapir = "0.1.2"
+Tapir = "0.2.4"
 Test = "1"
 Tracker = "0.2.33"
 Zygote = "0.6.69"

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/DifferentiationInterfaceTapirExt.jl

@@ -14,7 +14,11 @@ using Tapir:
     tangent_type,
     value_and_pullback!!,
     zero_codual,
-    zero_tangent
+    zero_tangent,
+    NoRData,
+    fdata,
+    rdata,
+    __value_and_pullback!!
 
 DI.check_available(::AutoTapir) = true
 

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/onearg.jl

@@ -5,25 +5,23 @@ end
 
 function DI.prepare_pullback(f, ::AutoTapir, x, dy)
     y = f(x)
-    rrule = build_rrule(f, x)
-    return TapirOneArgPullbackExtras(y, rrule)
+    return TapirOneArgPullbackExtras(y, build_rrule(f, x))
 end
 
 function DI.value_and_pullback(
     f, ::AutoTapir, x, dy, extras::TapirOneArgPullbackExtras{Y}
 ) where {Y}
     dy_righttype = convert(tangent_type(Y), dy)
-    new_y, (new_df, new_dx) = value_and_pullback!!(extras.rrule, dy_righttype, f, x)
+    new_y, (_, 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::TapirOneArgPullbackExtras{Y}
 ) where {Y}
     dy_righttype = convert(tangent_type(Y), dy)
-    dx_righttype = convert(tangent_type(typeof(x)), dx)
-    dx_righttype = set_to_zero!!(dx_righttype)
-    y, (new_df, new_dx) = value_and_pullback!!(
+    dx_righttype = set_to_zero!!(convert(tangent_type(typeof(x)), dx))
+    y, (_, new_dx) = __value_and_pullback!!(
         extras.rrule, dy_righttype, zero_codual(f), CoDual(x, dx_righttype)
     )
     return y, copyto!(dx, new_dx)

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/twoarg.jl

@@ -29,11 +29,13 @@ function DI.value_and_pullback(f!, y, ::AutoTapir, x, dy, extras::TapirTwoArgPul
 
     # Run the forwards-pass.
     out, pb!! = extras.rrule(
-        CoDual(f!, df!), CoDual(y_copy, dy_righttype), CoDual(x, dx_righttype)
+        CoDual(f!, fdata(df!)),
+        CoDual(y_copy, fdata(dy_righttype)),
+        CoDual(x, fdata(dx_righttype)),
     )
 
     # Verify that the output is non-differentiable.
-    @assert tangent(out) == NoTangent()
+    @assert primal(out) === nothing
 
     # Set the cotangent of `y` to be equal to the requested value.
     dy_righttype = increment!!(dy_righttype, dy_righttype_backup)
@@ -42,7 +44,7 @@ function DI.value_and_pullback(f!, y, ::AutoTapir, x, dy, extras::TapirTwoArgPul
     y = copyto!(y, y_copy)
 
     # Run the reverse-pass.
-    _, _, new_dx = pb!!(NoTangent(), df!, dy_righttype, dx_righttype)
+    _, _, new_dx = pb!!(NoRData())
 
-    return y, new_dx
+    return y, tangent(fdata(dx_righttype), new_dx)
 end