Ensemble simulations produce solutions with different dimensionality between CPU and GPU (vec's the solution)

Solving a matrix differential equation yields a matrix solution (as expected) on a CPU, but on a GPU a vector solution is returned.

Suppose that the following is executed on the REPL:

using LinearAlgebra, DifferentialEquations, DiffEqGPU

dim = 2;
t0 = 0.0f0;
tf = 1.0f0;
steps = 11;
jobs = Float32[-1.0 1.0 1.0 1.0 0.0 0.0; -2.0 1.0 1.0 1.0 0.0 0.0];

tstep = (tf-t0)/(steps-1)

function du!(du, u, p, t)
    @inbounds begin
        du[1,1] = 1;#lower_d[1,1];
        du[1,2] = 1;#lower_d[1,2];
        du[2,1] = 1;#lower_d[2,1];
        du[2,2] = 1;#lower_d[2,2];
    end
    nothing
end

# Create the initial state
u0 = zeros(Complex{Float32},dim,dim);
u0[1,1] = 1;

prob = ODEProblem{true}(du!,u0,(t0,tf),jobs[1,:]);

# Define how the problem is modified for each job
# i defines the trajectory number in the ensembleproblem
# repeat is true if the problem is being repeated
function prob_func(prob,i,repeat)
    return remake(prob,p=jobs[i,:])
end

ensemble_prob = EnsembleProblem(prob,
                                prob_func=prob_func,
                                safetycopy=false)

The system is then solved on the CPU with:

sol = solve(ensemble_prob,
                   Tsit5(),
                   trajectories=size(jobs,1),
                   saveat=tstep,
                   dense=false,
                   save_everystep=false,
                   abstol=1e-10,
                   reltol=1e-7,
                   maxiters=1e5);

Inspecting an element of the solution yields:

julia> sol.u[1]
retcode: Success
Interpolation: 1st order linear
t: 11-element Array{Float32,1}:
 0.0
 0.1
 0.2
 0.3
 0.4
 0.5
 0.6
 0.7
 0.8
 0.9
 1.0
u: 11-element Array{Array{Complex{Float32},2},1}:
 [1.0f0 + 0.0f0im 0.0f0 + 0.0f0im; 0.0f0 + 0.0f0im 0.0f0 + 0.0f0im]
 [1.1000003f0 + 0.0f0im 0.1f0 + 0.0f0im; 0.1f0 + 0.0f0im 0.1f0 + 0.0f0im]
 [1.2000004f0 + 0.0f0im 0.2f0 + 0.0f0im; 0.2f0 + 0.0f0im 0.2f0 + 0.0f0im]
 [1.3000004f0 + 0.0f0im 0.29999998f0 + 0.0f0im; 0.29999998f0 + 0.0f0im 0.29999998f0 + 0.0f0im]
 [1.4000005f0 + 0.0f0im 0.40000004f0 + 0.0f0im; 0.40000004f0 + 0.0f0im 0.40000004f0 + 0.0f0im]
 [1.5000004f0 + 0.0f0im 0.5f0 + 0.0f0im; 0.5f0 + 0.0f0im 0.5f0 + 0.0f0im]
 [1.6000003f0 + 0.0f0im 0.6f0 + 0.0f0im; 0.6f0 + 0.0f0im 0.6f0 + 0.0f0im]
 [1.7000003f0 + 0.0f0im 0.7f0 + 0.0f0im; 0.7f0 + 0.0f0im 0.7f0 + 0.0f0im]
 [1.8000002f0 + 0.0f0im 0.79999995f0 + 0.0f0im; 0.79999995f0 + 0.0f0im 0.79999995f0 + 0.0f0im]
 [1.9000002f0 + 0.0f0im 0.8999999f0 + 0.0f0im; 0.8999999f0 + 0.0f0im 0.8999999f0 + 0.0f0im]
 [2.0000002f0 + 0.0f0im 1.0f0 + 0.0f0im; 1.0f0 + 0.0f0im 1.0f0 + 0.0f0im]

u0 and du are matrices of Complex{Float32}, as are the elements of the solution array.

Running this on the GPU by changing ensemblealg yields:

julia> sol = solve(ensemble_prob,
                   Tsit5(), EnsembleGPUArray(),
                   trajectories=size(jobs,1),
                   saveat=tstep,
                   dense=false,
                   save_everystep=false,
                   abstol=1e-10,
                   reltol=1e-7,
                   maxiters=1e5);

julia> sol.u[1]
retcode: Success
Interpolation: 1st order linear
t: 11-element Array{Float32,1}:
 0.0
 0.1
 0.2
 0.3
 0.4
 0.5
 0.6
 0.7
 0.8
 0.9
 1.0
u: 11-element Array{SubArray{Complex{Float32},1,Array{Complex{Float32},2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true},1}:
 [1.0f0 + 0.0f0im, 0.0f0 + 0.0f0im]
 [1.1000003f0 + 0.0f0im, 0.1f0 + 0.0f0im]
 [1.2000004f0 + 0.0f0im, 0.2f0 + 0.0f0im]
 [1.3000004f0 + 0.0f0im, 0.29999998f0 + 0.0f0im]
 [1.4000005f0 + 0.0f0im, 0.40000004f0 + 0.0f0im]
 [1.5000004f0 + 0.0f0im, 0.5f0 + 0.0f0im]
 [1.6000003f0 + 0.0f0im, 0.6f0 + 0.0f0im]
 [1.7000003f0 + 0.0f0im, 0.7f0 + 0.0f0im]
 [1.8000002f0 + 0.0f0im, 0.79999995f0 + 0.0f0im]
 [1.9000002f0 + 0.0f0im, 0.8999999f0 + 0.0f0im]
 [2.0000002f0 + 0.0f0im, 1.0f0 + 0.0f0im]

The elements of the output seem to be an array of vectors, not matrices.

Reproduce the discrepancy from the provided REPL example, comparing the default CPU ensemble solve with EnsembleGPUArray(). Start at the EnsembleGPUArray entry point and trace how matrix states are stored and returned; done means GPU ensemble solutions preserve the matrix dimensionality shown by the CPU result.