Additional allocation to convert index array to same type as x/y

Author: gdalle

DifferentiationInterface/src/first_order/pullback.jl

@@ -168,9 +168,11 @@ function _pullback_via_pushforward(
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    dx = map(x, CartesianIndices(x)) do xj, j
+    ind = CartesianIndices(x)
+    T = typeof(similar(x, eltype(ind)))
+    dx = map(x, T(ind)) do xj, j
         t1 = pushforward(f, pushforward_prep, backend, x, (basis(x, j),), contexts...)
-        convert(eltype(x), dot(only(t1), dy))
+        dot(only(t1), dy)
     end
     return dx
 end
@@ -254,9 +256,11 @@ function _pullback_via_pushforward(
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    dx = map(x, CartesianIndices(x)) do xj, j  # preserve shape
+    ind = CartesianIndices(x)
+    T = typeof(similar(x, eltype(ind)))
+    dx = map(x, T(ind)) do xj, j  # preserve shape
         t1 = pushforward(f!, y, pushforward_prep, backend, x, (basis(x, j),), contexts...)
-        convert(eltype(x), dot(only(t1), dy))
+        dot(only(t1), dy)
     end
     return dx
 end

DifferentiationInterface/src/first_order/pushforward.jl

@@ -171,9 +171,11 @@ function _pushforward_via_pullback(
     dx,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    dy = map(y, CartesianIndices(y)) do yi, i
+    ind = CartesianIndices(y)
+    T = typeof(similar(y, eltype(ind)))
+    dy = map(y, T(ind)) do yi, i
         t1 = pullback(f, pullback_prep, backend, x, (basis(y, i),), contexts...)
-        convert(eltype(y), dot(only(t1), dx))
+        dot(only(t1), dx)
     end
     return dy
 end
@@ -243,9 +245,11 @@ function _pushforward_via_pullback(
     dx,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    dy = map(y, CartesianIndices(y)) do yi, i  # preserve shape
+    ind = CartesianIndices(y)
+    T = typeof(similar(y, eltype(ind)))
+    dy = map(y, T(ind)) do yi, i  # preserve shape
         t1 = pullback(f!, y, pullback_prep, backend, x, (basis(y, i),), contexts...)
-        convert(eltype(y), dot(only(t1), dx))
+        dot(only(t1), dx)
     end
     return dy
 end