# Calculation of gradient not working for an ODEproblem that uses 2D or 3D arrays

Solving a set of 30 coupled pairs of 2 differential equations with u, du a 2x30 2D array results in

ERROR: DimensionMismatch(“arrays could not be broadcast to a common size; got a dimension with lengths 2 and 60”).

The problem is that internally the vector of dimension 2x30 is flattened into a 60 element vector. Same problem occurs for 3D arrays.

``````using Flux
using DifferentialEquations
using DiffEqSensitivity

function model(p)
N_oscillators = 30
u0 = repeat([0.0; 1.0], 1, N_oscillators) # size(u0) = (2, 30)

function du!(du, u, p, t)
W, b  = p # Parameters
dy  = @view du[1,:] # 30 elements
dy′ = @view du[2,:]
y = @view u[1,:]
y′= @view u[2,:]
@. dy′ = -y * W
@. dy  = y′ * b
end

output = solve(
ODEProblem(
du!,
u0,
(0.0, 10.0),
p,
jac = true,
abstol = 1e-12,
reltol = 1e-12),
Tsit5(),
jac = true,
saveat = collect(0:0.1:7),
)
return Array(output[1, :, :]) # only return y, not y′
end

p=[1.5, 0.1]
y = model(p)
loss() = sum(model(p))
@show gradient(() -> loss(), params(p)) # This fails with a Broadcast Error

``````

Package Version:

Julia Version 1.4.0
Commit b8e9a9ecc6 (2020-03-21 16:36 UTC)
Platform Info:
OS: Windows (x86_64-w64-mingw32)
CPU: Intel® Core™ i7-6820HQ CPU @ 2.70GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-8.0.1 (ORCJIT, skylake)

[587475ba] Flux v0.10.4
[0c46a032] DifferentialEquations v6.15.0
[41bf760c] DiffEqSensitivity v6.23.0

I’d be grateful of any insights into the solution.

Stacktrace:

[4] combine_axes at .julia\packages\DiffEqBase\uSSHl\src\diffeqfastbc.jl:23 [inlined]

[8] _vecjacobian! at C:\Users\user.julia\packages\DiffEqSensitivity\XQ1vJ\src\local_sensitivity\derivative_wrappers.jl:185 [inlined]

[9] #vecjacobian!#20 at C:\Users\user.julia\packages\DiffEqSensitivity\XQ1vJ\src\local_sensitivity\derivative_wrappers.jl:147 [inlined]

[10] vecjacobian! at C:\Users\user.julia\packages\DiffEqSensitivity\XQ1vJ\src\local_sensitivity\derivative_wrappers.jl:147 [inlined]