Adapt to latest Mooncake

Author: gdalle

DifferentiationInterface/CHANGELOG.md

@@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [0.7.1]
 
+### Feat
+
+- Use Mooncake's internal copy utilities ([#809])
+
 ### Fixed
 
 - Make basis work for `CuArray` ([#810])
@@ -40,6 +44,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [0.6.53]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.52...DifferentiationInterface-v0.6.53
 
 [#810]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/810
+[#809]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/809
 [#799]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/799
 [#795]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/795
 [#790]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/790

DifferentiationInterface/Project.toml

@@ -74,7 +74,7 @@ JET = "0.9"
 JLArrays = "0.2.0"
 JuliaFormatter = "1,2"
 LinearAlgebra = "1"
-Mooncake = "0.4.121"
+Mooncake = "0.4.122"
 Pkg = "1"
 PolyesterForwardDiff = "0.1.2"
 Random = "1"

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -12,7 +12,7 @@ using Mooncake:
     value_and_pullback!!,
     zero_tangent,
     _copy_output,
-    _copy_to_output!
+    _copy_to_output!!
 
 DI.check_available(::AutoMooncake) = true
 

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/onearg.jl

@@ -30,7 +30,7 @@ function DI.value_and_pullback(
 ) where {F,Y,C}
     DI.check_prep(f, prep, backend, x, ty, contexts...)
     dy = only(ty)
-    dy_righttype = dy isa tangent_type(Y) ? dy : _copy_to_output!(prep.dy_righttype, dy)
+    dy_righttype = dy isa tangent_type(Y) ? dy : _copy_to_output!!(prep.dy_righttype, dy)
     new_y, (_, new_dx) = value_and_pullback!!(
         prep.cache, dy_righttype, f, x, map(DI.unwrap, contexts)...
     )
@@ -47,7 +47,8 @@ function DI.value_and_pullback(
 ) where {F,Y,C}
     DI.check_prep(f, prep, backend, x, ty, contexts...)
     ys_and_tx = map(ty) do dy
-        dy_righttype = dy isa tangent_type(Y) ? dy : _copy_to_output!(prep.dy_righttype, dy)
+        dy_righttype =
+            dy isa tangent_type(Y) ? dy : _copy_to_output!!(prep.dy_righttype, dy)
         y, (_, new_dx) = value_and_pullback!!(
             prep.cache, dy_righttype, f, x, map(DI.unwrap, contexts)...
         )