feat: use Mooncake's copy utilities

Author: gdalle

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.7.0"
+version = "0.7.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -41,7 +41,9 @@ DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
 DifferentiationInterfaceGPUArraysCoreExt = "GPUArraysCore"
 DifferentiationInterfaceGTPSAExt = "GTPSA"
 DifferentiationInterfaceMooncakeExt = "Mooncake"
-DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"]
+DifferentiationInterfacePolyesterForwardDiffExt = [
+    "PolyesterForwardDiff", "ForwardDiff", "DiffResults"
+]
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
 DifferentiationInterfaceSparseArraysExt = "SparseArrays"
 DifferentiationInterfaceSparseConnectivityTracerExt = "SparseConnectivityTracer"
@@ -53,7 +55,7 @@ DifferentiationInterfaceZygoteExt = ["Zygote", "ForwardDiff"]
 
 [compat]
 Aqua = "0.8.12"
-ADTypes = "1.13.0"
+ADTypes = "1.15.0"
 ChainRulesCore = "1.23.0"
 ComponentArrays = "0.15.27"
 DataFrames = "1.7.0"
@@ -72,7 +74,7 @@ JET = "0.9"
 JLArrays = "0.2.0"
 JuliaFormatter = "1,2"
 LinearAlgebra = "1"
-Mooncake = "0.4.88"
+Mooncake = "0.4.121"
 Pkg = "1"
 PolyesterForwardDiff = "0.1.2"
 Random = "1"
@@ -121,4 +123,21 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ADTypes", "Aqua", "ComponentArrays", "DataFrames", "ExplicitImports", "JET", "JLArrays", "JuliaFormatter", "Pkg", "Random", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test"]
+test = [
+    "ADTypes",
+    "Aqua",
+    "ComponentArrays",
+    "DataFrames",
+    "ExplicitImports",
+    "JET",
+    "JLArrays",
+    "JuliaFormatter",
+    "Pkg",
+    "Random",
+    "SparseArrays",
+    "SparseConnectivityTracer",
+    "SparseMatrixColorings",
+    "StableRNGs",
+    "StaticArrays",
+    "Test",
+]

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -10,13 +10,12 @@ using Mooncake:
     tangent_type,
     value_and_gradient!!,
     value_and_pullback!!,
-    zero_tangent
+    zero_tangent,
+    _copy_output,
+    _copy_to_output!
 
 DI.check_available(::AutoMooncake) = true
 
-copyto!!(dst::Number, src::Number) = convert(typeof(dst), src)
-copyto!!(dst, src) = DI.ismutable_array(dst) ? copyto!(dst, src) : convert(typeof(dst), src)
-
 get_config(::AutoMooncake{Nothing}) = Config()
 get_config(backend::AutoMooncake{<:Config}) = backend.config
 

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/onearg.jl

@@ -30,11 +30,11 @@ 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 : copyto!!(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)...
     )
-    return new_y, (mycopy(new_dx),)
+    return new_y, (_copy_output(new_dx),)
 end
 
 function DI.value_and_pullback(
@@ -47,11 +47,11 @@ 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 : copyto!!(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)...
         )
-        y, mycopy(new_dx)
+        y, _copy_output(new_dx)
     end
     y = first(ys_and_tx[1])
     tx = last.(ys_and_tx)
@@ -126,7 +126,7 @@ function DI.value_and_gradient(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     y, (_, new_grad) = value_and_gradient!!(prep.cache, f, x, map(DI.unwrap, contexts)...)
-    return y, mycopy(new_grad)
+    return y, _copy_output(new_grad)
 end
 
 function DI.value_and_gradient!(

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/twoarg.jl

@@ -58,7 +58,7 @@ function DI.value_and_pullback(
         map(DI.unwrap, contexts)...,
     )
     copyto!(y, y_after)
-    return y, (mycopy(dx),)
+    return y, (_copy_output(dx),)
 end
 
 function DI.value_and_pullback(
@@ -83,7 +83,7 @@ function DI.value_and_pullback(
             map(DI.unwrap, contexts)...,
         )
         copyto!(y, y_after)
-        mycopy(dx)
+        _copy_output(dx)
     end
     return y, tx
 end