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

[BillyKalfus]

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.

[ChrisRackauckas]

Thanks for the report. Yes, it looks like this is an artifact of how the SPMD kernel is generated and it should get a reshape before changing back to the CPU.