fix: make wrong-mode pushforward/pullback return the correct array type

Author: gdalle

DifferentiationInterface/Project.toml

@@ -1,13 +1,14 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
-authors = ["Guillaume Dalle", "Adrian Hill"]
 version = "0.7.16"
+authors = ["Guillaume Dalle", "Adrian Hill"]
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [weakdeps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
@@ -38,7 +39,7 @@ DifferentiationInterfaceFastDifferentiationExt = "FastDifferentiation"
 DifferentiationInterfaceFiniteDiffExt = "FiniteDiff"
 DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences"
 DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
-DifferentiationInterfaceGPUArraysCoreExt = "GPUArraysCore"
+DifferentiationInterfaceGPUArraysCoreExt = ["GPUArraysCore", "Adapt"]
 DifferentiationInterfaceGTPSAExt = "GTPSA"
 DifferentiationInterfaceMooncakeExt = "Mooncake"
 DifferentiationInterfacePolyesterForwardDiffExt = [

DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl

@@ -1,5 +1,6 @@
 module DifferentiationInterfaceGPUArraysCoreExt
 
+using Adapt: adapt
 import DifferentiationInterface as DI
 using GPUArraysCore: @allowscalar, AbstractGPUArray
 
@@ -17,4 +18,10 @@ function DI.multibasis(a::AbstractGPUArray{T}, inds) where {T}
     return b
 end
 
+function DI.arroftup_to_tupofarr(
+        tx::AbstractArray{NTuple{B, T}}, x::AbstractGPUArray{T}
+    ) where {B, T}
+    return ntuple(b -> adapt(typeof(x), getindex.(tx, b)), Val(B))
+end
+
 end

DifferentiationInterface/src/first_order/pullback.jl

@@ -332,7 +332,7 @@ function _value_and_pullback_via_pushforward(
     tx = map(ty) do dy
         dot(a, dy)
     end
-    return y, arroftup_to_tupofarr(tx)
+    return y, arroftup_to_tupofarr(tx, x)
 end
 
 function _value_and_pullback_via_pushforward(
@@ -348,7 +348,7 @@ function _value_and_pullback_via_pushforward(
     tx = map(ty) do dy
         real(dot(a, dy)) + im * real(dot(b, dy))
     end
-    return y, arroftup_to_tupofarr(tx)
+    return y, arroftup_to_tupofarr(tx, x)
 end
 
 function _value_and_pullback_via_pushforward(
@@ -366,7 +366,7 @@ function _value_and_pullback_via_pushforward(
             dot(a, dy)
         end
     end
-    return y, arroftup_to_tupofarr(tx)
+    return y, arroftup_to_tupofarr(tx, x)
 end
 
 function _value_and_pullback_via_pushforward(
@@ -387,7 +387,7 @@ function _value_and_pullback_via_pushforward(
             real(dot(a, dy)) + im * real(dot(b, dy))
         end
     end
-    return y, arroftup_to_tupofarr(tx)
+    return y, arroftup_to_tupofarr(tx, x)
 end
 
 function value_and_pullback(
@@ -458,7 +458,7 @@ function _value_and_pullback_via_pushforward(
     tx = map(ty) do dy
         dot(a, dy)
     end
-    return y, arroftup_to_tupofarr(tx)
+    return y, arroftup_to_tupofarr(tx, x)
 end
 
 function _value_and_pullback_via_pushforward(
@@ -477,7 +477,7 @@ function _value_and_pullback_via_pushforward(
     tx = map(ty) do dy
         real(dot(a, dy)) + im * real(dot(b, dy))
     end
-    return y, arroftup_to_tupofarr(tx)
+    return y, arroftup_to_tupofarr(tx, x)
 end
 
 function _value_and_pullback_via_pushforward(
@@ -495,7 +495,7 @@ function _value_and_pullback_via_pushforward(
             dot(a, dy)
         end
     end
-    return y, arroftup_to_tupofarr(tx)
+    return y, arroftup_to_tupofarr(tx, x)
 end
 
 function _value_and_pullback_via_pushforward(
@@ -518,7 +518,7 @@ function _value_and_pullback_via_pushforward(
             real(dot(a, dy)) + im * real(dot(b, dy))
         end
     end
-    return y, arroftup_to_tupofarr(tx)
+    return y, arroftup_to_tupofarr(tx, x)
 end
 
 function value_and_pullback(

DifferentiationInterface/src/first_order/pushforward.jl

@@ -331,7 +331,7 @@ function _value_and_pushforward_via_pullback(
     ty = map(tx) do dx
         dot(a, dx)
     end
-    return y, arroftup_to_tupofarr(ty)
+    return y, arroftup_to_tupofarr(ty, y)
 end
 
 function _value_and_pushforward_via_pullback(
@@ -348,7 +348,7 @@ function _value_and_pushforward_via_pullback(
     ty = map(tx) do dx
         real(dot(a, dx)) + im * real(dot(b, dx))
     end
-    return y, arroftup_to_tupofarr(ty)
+    return y, arroftup_to_tupofarr(ty, y)
 end
 
 function _value_and_pushforward_via_pullback(
@@ -367,7 +367,7 @@ function _value_and_pushforward_via_pullback(
             dot(a, dx)
         end
     end
-    return y, arroftup_to_tupofarr(ty)
+    return y, arroftup_to_tupofarr(ty, y)
 end
 
 function _value_and_pushforward_via_pullback(
@@ -387,7 +387,7 @@ function _value_and_pushforward_via_pullback(
             real(dot(a, dx)) + im * real(dot(b, dx))
         end
     end
-    return y, arroftup_to_tupofarr(ty)
+    return y, arroftup_to_tupofarr(ty, y)
 end
 
 function value_and_pushforward(
@@ -460,7 +460,7 @@ function _value_and_pushforward_via_pullback(
             dot(a, dx)
         end
     end
-    return y, arroftup_to_tupofarr(ty)
+    return y, arroftup_to_tupofarr(ty, y)
 end
 
 function _value_and_pushforward_via_pullback(
@@ -481,7 +481,7 @@ function _value_and_pushforward_via_pullback(
             real(dot(a, dx)) + im * real(dot(b, dx))
         end
     end
-    return y, arroftup_to_tupofarr(ty)
+    return y, arroftup_to_tupofarr(ty, y)
 end
 
 function value_and_pushforward(

DifferentiationInterface/src/utils/linalg.jl

@@ -36,5 +36,13 @@ get_pattern(M::AbstractMatrix) = trues(size(M))
 
 onlysecond((a, b)) = (a, only(b))
 
-arroftup_to_tupofarr(x::NTuple) = x
-arroftup_to_tupofarr(x::AbstractArray{<:NTuple{B}}) where {B} = ntuple(b -> getindex.(x, b), Val(B))
+"""
+    arroftup_to_tupofarr(tx, x)
+
+Convert an array of tuples `tx` into a tuple of arrays, while respecting the array type of the primal `x`.
+"""
+arroftup_to_tupofarr(tx::NTuple{B, T}, x::T) where {B, T} = tx
+
+function arroftup_to_tupofarr(tx::AbstractArray{NTuple{B, T}}, x::AbstractArray{T}) where {B, T}
+    return ntuple(b -> similar(x) .= getindex.(tx, b), Val(B))
+end