fix: respect array type in wrong-mode pushforward/pullback

Author: gdalle

DifferentiationInterface/src/first_order/pullback.jl

@@ -167,8 +167,8 @@ function _pullback_via_pushforward(
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    t1 = pushforward(f, pushforward_prep, backend, x, (one(x),), contexts...)
-    dx = dot(dy, only(t1))
+    t = pushforward(f, pushforward_prep, backend, x, (one(x),), contexts...)
+    dx = dot(dy, only(t))
     return dx
 end
 
@@ -180,9 +180,10 @@ function _pullback_via_pushforward(
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    dx = map(CartesianIndices(x)) do j
-        t1 = pushforward(f, pushforward_prep, backend, x, (basis(backend, x, j),), contexts...)
-        dot(dy, only(t1))
+    dx = map(x, CartesianIndices(x)) do xj, j
+        bj = basis(backend, x, j)
+        tj = pushforward(f, pushforward_prep, backend, x, (bj,), contexts...)
+        dot(dy, only(tj))
     end
     return dx
 end
@@ -252,8 +253,8 @@ function _pullback_via_pushforward(
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    t1 = pushforward(f!, y, pushforward_prep, backend, x, (one(x),), contexts...)
-    dx = dot(dy, only(t1))
+    t = pushforward(f!, y, pushforward_prep, backend, x, (one(x),), contexts...)
+    dx = dot(dy, only(t))
     return dx
 end
 
@@ -266,11 +267,10 @@ function _pullback_via_pushforward(
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    dx = map(CartesianIndices(x)) do j  # preserve shape
-        t1 = pushforward(
-            f!, y, pushforward_prep, backend, x, (basis(backend, x, j),), contexts...
-        )
-        dot(dy, only(t1))
+    dx = map(x, CartesianIndices(x)) do xj, j  # preserve shape
+        bj = basis(backend, x, j)
+        tj = pushforward(f!, y, pushforward_prep, backend, x, (bj,), contexts...)
+        dot(dy, only(tj))
     end
     return dx
 end

DifferentiationInterface/src/first_order/pushforward.jl

@@ -175,8 +175,8 @@ function _pushforward_via_pullback(
     dx,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    t1 = pullback(f, pullback_prep, backend, x, (one(y),), contexts...)
-    dy = dot(dx, only(t1))
+    t = pullback(f, pullback_prep, backend, x, (one(y),), contexts...)
+    dy = dot(dx, only(t))
     return dy
 end
 
@@ -189,9 +189,10 @@ function _pushforward_via_pullback(
     dx,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    dy = map(CartesianIndices(y)) do i
-        t1 = pullback(f, pullback_prep, backend, x, (basis(backend, y, i),), contexts...)
-        dot(dx, only(t1))
+    dy = map(y, CartesianIndices(y)) do yi, i
+        bi = basis(backend, y, i)
+        ti = pullback(f, pullback_prep, backend, x, (bi,), contexts...)
+        dot(dx, only(ti))
     end
     return dy
 end
@@ -261,9 +262,10 @@ function _pushforward_via_pullback(
     dx,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    dy = map(CartesianIndices(y)) do i  # preserve shape
-        t1 = pullback(f!, y, pullback_prep, backend, x, (basis(backend, y, i),), contexts...)
-        dot(dx, only(t1))
+    dy = map(y, CartesianIndices(y)) do yi, i  # preserve shape
+        bi = basis(backend, y, i)
+        ti = pullback(f!, y, pullback_prep, backend, x, (bi,), contexts...)
+        dot(dx, only(ti))
     end
     return dy
 end