From ff529cb66c686f6357f9fe1c78648d70ae0feefc Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Thu, 27 Jun 2024 00:03:13 +0200
Subject: [PATCH] Improve type stability of Jacobians and Hessian, fix test
 scenarios (#337)

* Ensure type stability of test scenarios in 1.11

* Preallocate results

* Typo

* Fix GPU scenario

* Remporarily disable tests on 1.11

* I said disable them
---
 .github/workflows/Test.yml                    |  4 +-
 .../src/first_order/jacobian.jl               | 72 +++++++++++--------
 .../src/second_order/hessian.jl               | 39 +++++-----
 .../src/sparse/hessian.jl                     | 38 +++++++---
 .../src/sparse/jacobian.jl                    | 72 ++++++++++++-------
 DifferentiationInterface/src/utils/batch.jl   |  2 +
 .../src/scenarios/default.jl                  | 16 ++---
 7 files changed, 150 insertions(+), 93 deletions(-)

diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml
index 6027c23c1..4eef3524a 100644
--- a/.github/workflows/Test.yml
+++ b/.github/workflows/Test.yml
@@ -29,7 +29,7 @@ jobs:
         version:
           - '1'
           - '1.6'
-          - '~1.11.0-0'
+          # - '~1.11.0-0'
         group:
           - Formalities
           - Internals
@@ -118,7 +118,7 @@ jobs:
         version:
           - '1'
           - '1.6'
-          - '~1.11.0-0'
+          # - '~1.11.0-0'
         group:
           - Formalities
           - Zero
diff --git a/DifferentiationInterface/src/first_order/jacobian.jl b/DifferentiationInterface/src/first_order/jacobian.jl
index 3fe857068..513633cc2 100644
--- a/DifferentiationInterface/src/first_order/jacobian.jl
+++ b/DifferentiationInterface/src/first_order/jacobian.jl
@@ -55,16 +55,18 @@ abstract type JacobianExtras <: Extras end
 
 struct NoJacobianExtras <: JacobianExtras end
 
-struct PushforwardJacobianExtras{B,D,E<:PushforwardExtras,Y} <: JacobianExtras
+struct PushforwardJacobianExtras{B,D,R,E<:PushforwardExtras} <: JacobianExtras
     batched_seeds::Vector{Batch{B,D}}
+    batched_results::Vector{Batch{B,R}}
     pushforward_batched_extras::E
-    y_example::Y
+    N::Int
 end
 
-struct PullbackJacobianExtras{B,D,E<:PullbackExtras,Y} <: JacobianExtras
+struct PullbackJacobianExtras{B,D,R,E<:PullbackExtras} <: JacobianExtras
     batched_seeds::Vector{Batch{B,D}}
+    batched_results::Vector{Batch{B,R}}
     pullback_batched_extras::E
-    y_example::Y
+    M::Int
 end
 
 function prepare_jacobian(f::F, backend::AbstractADType, x) where {F}
@@ -85,14 +87,15 @@ function prepare_jacobian_aux(f_or_f!y::FY, backend, x, y, ::PushforwardFast) wh
             ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % N], Val(B)) for
             a in 1:div(N, B, RoundUp)
         ])
+    batched_results = Batch.([ntuple(b -> similar(y), Val(B)) for _ in batched_seeds])
     pushforward_batched_extras = prepare_pushforward_batched(
         f_or_f!y..., backend, x, batched_seeds[1]
     )
-    D = eltype(seeds)
+    D = eltype(batched_seeds[1])
+    R = eltype(batched_results[1])
     E = typeof(pushforward_batched_extras)
-    Y = typeof(y)
-    return PushforwardJacobianExtras{B,D,E,Y}(
-        batched_seeds, pushforward_batched_extras, copy(y)
+    return PushforwardJacobianExtras{B,D,R,E}(
+        batched_seeds, batched_results, pushforward_batched_extras, N
     )
 end
 
@@ -105,13 +108,16 @@ function prepare_jacobian_aux(f_or_f!y::FY, backend, x, y, ::PushforwardSlow) wh
             ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % M], Val(B)) for
             a in 1:div(M, B, RoundUp)
         ])
+    batched_results = Batch.([ntuple(b -> similar(x), Val(B)) for _ in batched_seeds])
     pullback_batched_extras = prepare_pullback_batched(
         f_or_f!y..., backend, x, batched_seeds[1]
     )
-    D = eltype(seeds)
+    D = eltype(batched_seeds[1])
+    R = eltype(batched_results[1])
     E = typeof(pullback_batched_extras)
-    Y = typeof(y)
-    return PullbackJacobianExtras{B,D,E,Y}(batched_seeds, pullback_batched_extras, copy(y))
+    return PullbackJacobianExtras{B,D,R,E}(
+        batched_seeds, batched_results, pullback_batched_extras, M
+    )
 end
 
 ## One argument
@@ -209,8 +215,7 @@ end
 function jacobian_aux(
     f_or_f!y::FY, backend, x::AbstractArray, extras::PushforwardJacobianExtras{B}
 ) where {FY,B}
-    @compat (; batched_seeds, pushforward_batched_extras, y_example) = extras
-    N = length(x)
+    @compat (; batched_seeds, pushforward_batched_extras, N) = extras
 
     pushforward_batched_extras_same = prepare_pushforward_batched_same_point(
         f_or_f!y..., backend, x, batched_seeds[1], pushforward_batched_extras
@@ -233,8 +238,7 @@ end
 function jacobian_aux(
     f_or_f!y::FY, backend, x::AbstractArray, extras::PullbackJacobianExtras{B}
 ) where {FY,B}
-    @compat (; batched_seeds, pullback_batched_extras, y_example) = extras
-    M = length(y_example)
+    @compat (; batched_seeds, pullback_batched_extras, M) = extras
 
     pullback_batched_extras_same = prepare_pullback_batched_same_point(
         f_or_f!y..., backend, x, batched_seeds[1], extras.pullback_batched_extras
@@ -261,20 +265,16 @@ function jacobian_aux!(
     x::AbstractArray,
     extras::PushforwardJacobianExtras{B},
 ) where {FY,B}
-    @compat (; batched_seeds, pushforward_batched_extras, y_example) = extras
-    N = length(x)
+    @compat (; batched_seeds, batched_results, pushforward_batched_extras, N) = extras
 
     pushforward_batched_extras_same = prepare_pushforward_batched_same_point(
         f_or_f!y..., backend, x, batched_seeds[1], pushforward_batched_extras
     )
 
-    for a in eachindex(batched_seeds)
-        dy_batch_elements = ntuple(Val(B)) do b
-            reshape(view(jac, :, 1 + ((a - 1) * B + (b - 1)) % N), size(y_example))
-        end
+    for a in eachindex(batched_seeds, batched_results)
         pushforward_batched!(
             f_or_f!y...,
-            Batch(dy_batch_elements),
+            batched_results[a],
             backend,
             x,
             batched_seeds[a],
@@ -282,6 +282,15 @@ function jacobian_aux!(
         )
     end
 
+    for a in eachindex(batched_results)
+        for b in eachindex(batched_results[a].elements)
+            copyto!(
+                view(jac, :, 1 + ((a - 1) * B + (b - 1)) % N),
+                vec(batched_results[a].elements[b]),
+            )
+        end
+    end
+
     return jac
 end
 
@@ -292,20 +301,16 @@ function jacobian_aux!(
     x::AbstractArray,
     extras::PullbackJacobianExtras{B},
 ) where {FY,B}
-    @compat (; batched_seeds, pullback_batched_extras, y_example) = extras
-    M = length(y_example)
+    @compat (; batched_seeds, batched_results, pullback_batched_extras, M) = extras
 
     pullback_batched_extras_same = prepare_pullback_batched_same_point(
         f_or_f!y..., backend, x, batched_seeds[1], extras.pullback_batched_extras
     )
 
-    for a in eachindex(batched_seeds)
-        dx_batch_elements = ntuple(Val(B)) do b
-            reshape(view(jac, 1 + ((a - 1) * B + (b - 1)) % M, :), size(x))
-        end
+    for a in eachindex(batched_seeds, batched_results)
         pullback_batched!(
             f_or_f!y...,
-            Batch(dx_batch_elements),
+            batched_results[a],
             backend,
             x,
             batched_seeds[a],
@@ -313,5 +318,14 @@ function jacobian_aux!(
         )
     end
 
+    for a in eachindex(batched_results)
+        for b in eachindex(batched_results[a].elements)
+            copyto!(
+                view(jac, 1 + ((a - 1) * B + (b - 1)) % M, :),
+                vec(batched_results[a].elements[b]),
+            )
+        end
+    end
+
     return jac
 end
diff --git a/DifferentiationInterface/src/second_order/hessian.jl b/DifferentiationInterface/src/second_order/hessian.jl
index 3f68d7dfb..d879e7169 100644
--- a/DifferentiationInterface/src/second_order/hessian.jl
+++ b/DifferentiationInterface/src/second_order/hessian.jl
@@ -49,10 +49,12 @@ abstract type HessianExtras <: Extras end
 
 struct NoHessianExtras <: HessianExtras end
 
-struct HVPGradientHessianExtras{B,D,E2<:HVPExtras,E1<:GradientExtras} <: HessianExtras
+struct HVPGradientHessianExtras{B,D,R,E2<:HVPExtras,E1<:GradientExtras} <: HessianExtras
     batched_seeds::Vector{Batch{B,D}}
+    batched_results::Vector{Batch{B,R}}
     hvp_batched_extras::E2
     gradient_extras::E1
+    N::Int
 end
 
 function prepare_hessian(f::F, backend::AbstractADType, x) where {F}
@@ -64,12 +66,14 @@ function prepare_hessian(f::F, backend::AbstractADType, x) where {F}
             ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % N], Val(B)) for
             a in 1:div(N, B, RoundUp)
         ])
+    batched_results = Batch.([ntuple(b -> similar(x), Val(B)) for _ in batched_seeds])
     hvp_batched_extras = prepare_hvp_batched(f, backend, x, batched_seeds[1])
     gradient_extras = prepare_gradient(f, maybe_inner(backend), x)
-    D = eltype(seeds)
+    D = eltype(batched_seeds[1])
+    R = eltype(batched_results[1])
     E2, E1 = typeof(hvp_batched_extras), typeof(gradient_extras)
-    return HVPGradientHessianExtras{B,D,E2,E1}(
-        batched_seeds, hvp_batched_extras, gradient_extras
+    return HVPGradientHessianExtras{B,D,R,E2,E1}(
+        batched_seeds, batched_results, hvp_batched_extras, gradient_extras, N
     )
 end
 
@@ -100,8 +104,7 @@ end
 function hessian(
     f::F, backend::AbstractADType, x, extras::HVPGradientHessianExtras{B}
 ) where {F,B}
-    @compat (; batched_seeds, hvp_batched_extras) = extras
-    N = length(x)
+    @compat (; batched_seeds, hvp_batched_extras, N) = extras
 
     hvp_batched_extras_same = prepare_hvp_batched_same_point(
         f, backend, x, batched_seeds[1], hvp_batched_extras
@@ -122,27 +125,27 @@ end
 function hessian!(
     f::F, hess, backend::AbstractADType, x, extras::HVPGradientHessianExtras{B}
 ) where {F,B}
-    @compat (; batched_seeds, hvp_batched_extras) = extras
-    N = length(x)
+    @compat (; batched_seeds, batched_results, hvp_batched_extras, N) = extras
 
     hvp_batched_extras_same = prepare_hvp_batched_same_point(
         f, backend, x, batched_seeds[1], hvp_batched_extras
     )
 
-    for a in eachindex(batched_seeds)
-        dg_batch_elements = ntuple(Val(B)) do b
-            reshape(view(hess, :, 1 + ((a - 1) * B + (b - 1)) % N), size(x))
-        end
+    for a in eachindex(batched_seeds, batched_results)
         hvp_batched!(
-            f,
-            Batch(dg_batch_elements),
-            backend,
-            x,
-            batched_seeds[a],
-            hvp_batched_extras_same,
+            f, batched_results[a], backend, x, batched_seeds[a], hvp_batched_extras_same
         )
     end
 
+    for a in eachindex(batched_results)
+        for b in eachindex(batched_results[a].elements)
+            copyto!(
+                view(hess, :, 1 + ((a - 1) * B + (b - 1)) % N),
+                vec(batched_results[a].elements[b]),
+            )
+        end
+    end
+
     return hess
 end
 
diff --git a/DifferentiationInterface/src/sparse/hessian.jl b/DifferentiationInterface/src/sparse/hessian.jl
index 0c4ee8224..869aecb2b 100644
--- a/DifferentiationInterface/src/sparse/hessian.jl
+++ b/DifferentiationInterface/src/sparse/hessian.jl
@@ -1,11 +1,12 @@
 struct SparseHessianExtras{
-    B,S<:AbstractMatrix{Bool},C<:AbstractMatrix{<:Real},D,E2<:HVPExtras,E1<:GradientExtras
+    B,S<:AbstractMatrix{Bool},C<:AbstractMatrix{<:Real},D,R,E2<:HVPExtras,E1<:GradientExtras
 } <: HessianExtras
     sparsity::S
     colors::Vector{Int}
     groups::Vector{Vector{Int}}
     compressed::C
     batched_seeds::Vector{Batch{B,D}}
+    batched_results::Vector{Batch{B,R}}
     hvp_batched_extras::E2
     gradient_extras::E1
 end
@@ -16,16 +17,18 @@ function SparseHessianExtras{B}(;
     groups,
     compressed::C,
     batched_seeds::Vector{Batch{B,D}},
+    batched_results::Vector{Batch{B,R}},
     hvp_batched_extras::E2,
     gradient_extras::E1,
-) where {B,S,C,D,E2,E1}
+) where {B,S,C,D,R,E2,E1}
     @assert size(sparsity, 1) == size(sparsity, 2) == size(compressed, 1) == length(colors)
-    return SparseHessianExtras{B,S,C,D,E2,E1}(
+    return SparseHessianExtras{B,S,C,D,R,E2,E1}(
         sparsity,
         colors,
         groups,
         compressed,
         batched_seeds,
+        batched_results,
         hvp_batched_extras,
         gradient_extras,
     )
@@ -48,6 +51,7 @@ function prepare_hessian(f::F, backend::AutoSparse, x) where {F}
             ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % Ng], Val(B)) for
             a in 1:div(Ng, B, RoundUp)
         ])
+    batched_results = Batch.([ntuple(b -> similar(x), Val(B)) for _ in batched_seeds])
     hvp_batched_extras = prepare_hvp_batched(f, dense_backend, x, batched_seeds[1])
     gradient_extras = prepare_gradient(f, maybe_inner(dense_backend), x)
     return SparseHessianExtras{B}(;
@@ -56,6 +60,7 @@ function prepare_hessian(f::F, backend::AutoSparse, x) where {F}
         groups,
         compressed,
         batched_seeds,
+        batched_results,
         hvp_batched_extras,
         gradient_extras,
     )
@@ -86,8 +91,15 @@ end
 function hessian!(
     f::F, hess, backend::AutoSparse, x, extras::SparseHessianExtras{B}
 ) where {F,B}
-    @compat (; sparsity, compressed, colors, groups, batched_seeds, hvp_batched_extras) =
-        extras
+    @compat (;
+        sparsity,
+        compressed,
+        colors,
+        groups,
+        batched_seeds,
+        batched_results,
+        hvp_batched_extras,
+    ) = extras
     dense_backend = dense_ad(backend)
     Ng = length(groups)
 
@@ -95,13 +107,10 @@ function hessian!(
         f, dense_backend, x, batched_seeds[1], hvp_batched_extras
     )
 
-    for a in 1:div(Ng, B, RoundUp)
-        dg_batch_elements = ntuple(Val(B)) do b
-            reshape(view(compressed, :, 1 + ((a - 1) * B + (b - 1)) % Ng), size(x))
-        end
+    for a in eachindex(batched_seeds, batched_results)
         hvp_batched!(
             f,
-            Batch(dg_batch_elements),
+            batched_results[a],
             dense_backend,
             x,
             batched_seeds[a],
@@ -109,6 +118,15 @@ function hessian!(
         )
     end
 
+    for a in eachindex(batched_results)
+        for b in eachindex(batched_results[a].elements)
+            copyto!(
+                view(compressed, :, 1 + ((a - 1) * B + (b - 1)) % Ng),
+                vec(batched_results[a].elements[b]),
+            )
+        end
+    end
+
     decompress_symmetric!(hess, sparsity, compressed, colors)
     return hess
 end
diff --git a/DifferentiationInterface/src/sparse/jacobian.jl b/DifferentiationInterface/src/sparse/jacobian.jl
index eee873dc8..64c478e8d 100644
--- a/DifferentiationInterface/src/sparse/jacobian.jl
+++ b/DifferentiationInterface/src/sparse/jacobian.jl
@@ -3,27 +3,27 @@
 abstract type SparseJacobianExtras <: JacobianExtras end
 
 struct PushforwardSparseJacobianExtras{
-    B,S<:AbstractMatrix{Bool},C<:AbstractMatrix{<:Real},D,E<:PushforwardExtras,Y
+    B,S<:AbstractMatrix{Bool},C<:AbstractMatrix{<:Real},D,R,E<:PushforwardExtras
 } <: SparseJacobianExtras
     sparsity::S
     colors::Vector{Int}
     groups::Vector{Vector{Int}}
     compressed::C
     batched_seeds::Vector{Batch{B,D}}
+    batched_results::Vector{Batch{B,R}}
     pushforward_batched_extras::E
-    y_example::Y
 end
 
 struct PullbackSparseJacobianExtras{
-    B,S<:AbstractMatrix{Bool},C<:AbstractMatrix{<:Real},D,E<:PullbackExtras,Y
+    B,S<:AbstractMatrix{Bool},C<:AbstractMatrix{<:Real},D,R,E<:PullbackExtras
 } <: SparseJacobianExtras
     sparsity::S
     colors::Vector{Int}
     groups::Vector{Vector{Int}}
     compressed::C
     batched_seeds::Vector{Batch{B,D}}
+    batched_results::Vector{Batch{B,R}}
     pullback_batched_extras::E
-    y_example::Y
 end
 
 function PushforwardSparseJacobianExtras{B}(;
@@ -32,19 +32,19 @@ function PushforwardSparseJacobianExtras{B}(;
     groups,
     compressed::C,
     batched_seeds::Vector{Batch{B,D}},
+    batched_results::Vector{Batch{B,R}},
     pushforward_batched_extras::E,
-    y_example::Y,
-) where {B,S,C,D,E,Y}
+) where {B,S,C,D,R,E}
     @assert size(sparsity, 1) == size(compressed, 1)
     @assert size(sparsity, 2) == length(colors)
-    return PushforwardSparseJacobianExtras{B,S,C,D,E,Y}(
+    return PushforwardSparseJacobianExtras{B,S,C,D,R,E}(
         sparsity,
         colors,
         groups,
         compressed,
         batched_seeds,
+        batched_results,
         pushforward_batched_extras,
-        y_example,
     )
 end
 
@@ -54,19 +54,19 @@ function PullbackSparseJacobianExtras{B}(;
     groups,
     compressed::C,
     batched_seeds::Vector{Batch{B,D}},
+    batched_results::Vector{Batch{B,R}},
     pullback_batched_extras::E,
-    y_example::Y,
-) where {B,S,C,D,E,Y}
+) where {B,S,C,D,R,E}
     @assert size(sparsity, 2) == size(compressed, 2)
     @assert size(sparsity, 1) == length(colors)
-    return PullbackSparseJacobianExtras{B,S,C,D,E,Y}(
+    return PullbackSparseJacobianExtras{B,S,C,D,R,E}(
         sparsity,
         colors,
         groups,
         compressed,
         batched_seeds,
+        batched_results,
         pullback_batched_extras,
-        y_example,
     )
 end
 
@@ -100,6 +100,7 @@ function prepare_sparse_jacobian_aux(
             ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % Ng], Val(B)) for
             a in 1:div(Ng, B, RoundUp)
         ])
+    batched_results = Batch.([ntuple(b -> similar(y), Val(B)) for _ in batched_seeds])
     pushforward_batched_extras = prepare_pushforward_batched(
         f_or_f!y..., dense_backend, x, batched_seeds[1]
     )
@@ -109,8 +110,8 @@ function prepare_sparse_jacobian_aux(
         groups,
         compressed,
         batched_seeds,
+        batched_results,
         pushforward_batched_extras,
-        y_example=copy(y),
     )
 end
 
@@ -131,6 +132,7 @@ function prepare_sparse_jacobian_aux(
             ntuple(b -> seeds[1 + ((a - 1) * B + (b - 1)) % Ng], Val(B)) for
             a in 1:div(Ng, B, RoundUp)
         ])
+    batched_results = Batch.([ntuple(b -> similar(x), Val(B)) for _ in batched_seeds])
     pullback_batched_extras = prepare_pullback_batched(
         f_or_f!y..., dense_backend, x, batched_seeds[1]
     )
@@ -140,8 +142,8 @@ function prepare_sparse_jacobian_aux(
         groups,
         compressed,
         batched_seeds,
+        batched_results,
         pullback_batched_extras,
-        y_example=copy(y),
     )
 end
 
@@ -270,8 +272,8 @@ function sparse_jacobian_aux!(
         colors,
         groups,
         batched_seeds,
+        batched_results,
         pushforward_batched_extras,
-        y_example,
     ) = extras
     dense_backend = dense_ad(backend)
     Ng = length(groups)
@@ -280,13 +282,10 @@ function sparse_jacobian_aux!(
         f_or_f!y..., dense_backend, x, batched_seeds[1], pushforward_batched_extras
     )
 
-    for a in eachindex(batched_seeds)
-        dy_batch_elements = ntuple(Val(B)) do b
-            reshape(view(compressed, :, 1 + ((a - 1) * B + (b - 1)) % Ng), size(y_example))
-        end
+    for a in eachindex(batched_seeds, batched_results)
         pushforward_batched!(
             f_or_f!y...,
-            Batch(dy_batch_elements),
+            batched_results[a],
             dense_backend,
             x,
             batched_seeds[a],
@@ -294,6 +293,15 @@ function sparse_jacobian_aux!(
         )
     end
 
+    for a in eachindex(batched_results)
+        for b in eachindex(batched_results[a].elements)
+            copyto!(
+                view(compressed, :, 1 + ((a - 1) * B + (b - 1)) % Ng),
+                vec(batched_results[a].elements[b]),
+            )
+        end
+    end
+
     decompress_columns!(jac, sparsity, compressed, colors)
     return jac
 end
@@ -302,7 +310,13 @@ function sparse_jacobian_aux!(
     f_or_f!y::FY, jac, backend::AutoSparse, x, extras::PullbackSparseJacobianExtras{B}
 ) where {FY,B}
     @compat (;
-        sparsity, compressed, colors, groups, batched_seeds, pullback_batched_extras
+        sparsity,
+        compressed,
+        colors,
+        groups,
+        batched_seeds,
+        batched_results,
+        pullback_batched_extras,
     ) = extras
     dense_backend = dense_ad(backend)
     Ng = length(groups)
@@ -311,13 +325,10 @@ function sparse_jacobian_aux!(
         f_or_f!y..., dense_backend, x, batched_seeds[1], pullback_batched_extras
     )
 
-    for a in eachindex(batched_seeds)
-        dx_batch_elements = ntuple(Val(B)) do b
-            reshape(view(compressed, 1 + ((a - 1) * B + (b - 1)) % Ng, :), size(x))
-        end
+    for a in eachindex(batched_seeds, batched_results)
         pullback_batched!(
             f_or_f!y...,
-            Batch(dx_batch_elements),
+            batched_results[a],
             dense_backend,
             x,
             batched_seeds[a],
@@ -325,6 +336,15 @@ function sparse_jacobian_aux!(
         )
     end
 
+    for a in eachindex(batched_results)
+        for b in eachindex(batched_results[a].elements)
+            copyto!(
+                view(compressed, 1 + ((a - 1) * B + (b - 1)) % Ng, :),
+                vec(batched_results[a].elements[b]),
+            )
+        end
+    end
+
     decompress_rows!(jac, sparsity, compressed, colors)
     return jac
 end
diff --git a/DifferentiationInterface/src/utils/batch.jl b/DifferentiationInterface/src/utils/batch.jl
index 7e2381818..d3ad56204 100644
--- a/DifferentiationInterface/src/utils/batch.jl
+++ b/DifferentiationInterface/src/utils/batch.jl
@@ -23,6 +23,8 @@ struct Batch{B,T}
     Batch(elements::NTuple) = new{length(elements),eltype(elements)}(elements)
 end
 
+Base.eltype(::Batch{B,T}) where {B,T} = T
+
 Base.:(==)(b1::Batch{B}, b2::Batch{B}) where {B} = b1.elements == b2.elements
 
 function Base.isapprox(b1::Batch{B}, b2::Batch{B}; kwargs...) where {B}
diff --git a/DifferentiationInterfaceTest/src/scenarios/default.jl b/DifferentiationInterfaceTest/src/scenarios/default.jl
index 66bea5671..7f620373d 100644
--- a/DifferentiationInterfaceTest/src/scenarios/default.jl
+++ b/DifferentiationInterfaceTest/src/scenarios/default.jl
@@ -38,8 +38,8 @@ end
 
 ## Number to array
 
-multiplicator(::Type{V}) where {V<:AbstractVector} = V(float.(1:6))
-multiplicator(::Type{M}) where {M<:AbstractMatrix} = M(float.(reshape(1:6, 2, 3)))
+multiplicator(::Type{V}) where {V<:AbstractVector} = convert(V, float.(1:6))
+multiplicator(::Type{M}) where {M<:AbstractMatrix} = convert(M, float.(reshape(1:6, 2, 3)))
 
 function num_to_arr(x::Number, ::Type{A}) where {A<:AbstractArray}
     a = multiplicator(A)
@@ -170,8 +170,8 @@ function arr_to_num_aux_no_linalg(x; α, β)
     return s
 end
 
-function arr_to_num_aux_gradient(x; α, β)
-    x = Array(x)  # GPU arrays don't like indexing
+function arr_to_num_aux_gradient(x0; α, β)
+    x = Array(x0)  # GPU arrays don't like indexing
     g = similar(x)
     for k in eachindex(g, x)
         g[k] = (
@@ -180,11 +180,11 @@ function arr_to_num_aux_gradient(x; α, β)
             (α + β) * x[k]^(α + β - 1)
         )
     end
-    return g
+    return convert(typeof(x0), g)
 end
 
-function arr_to_num_aux_hessian(x; α, β)
-    x = Array(x)  # GPU arrays don't like indexing
+function arr_to_num_aux_hessian(x0; α, β)
+    x = Array(x0)  # GPU arrays don't like indexing
     H = similar(x, length(x), length(x))
     for k in axes(H, 1), l in axes(H, 2)
         if k == l
@@ -197,7 +197,7 @@ function arr_to_num_aux_hessian(x; α, β)
             H[k, l] = α * β * (x[k]^(α - 1) * x[l]^(β - 1) + x[k]^(β - 1) * x[l]^(α - 1))
         end
     end
-    return H
+    return convert(typeof(similar(x0, length(x0), length(x0))), H)
 end
 
 const DEFAULT_α = 4