Pre-allocating arrays for parallel multiple shooting Neural ODE training in Flux

Specific Domains Machine Learning
, ,
1

Hi, I would like to pre-allocate an array of matrices which I would then in-place mutate in my multiple shooting training in Flux, optimized via Optimization (see skeleton code below). However, I don’t know what type should be used when pre-allocating it as Flux outputs some complex ForwardDiff.Dual{...} and don’t know how to get such type information.

A rough idea for the code that I am working with:

function loss(p, X, Y, NUM_SHOTS)

    pre_alloc_sols = [Array{typeToInstantiate}(undef, ncols, nrows) for i in 1:NUM_SHOTS

    for i in 1:NUM_SHOTS
        prob = DifferentialEquations.ODEProblem(...)
        sol  = DifferentialEquations.solve(prob, ...)

        pre_alloc_sols[i] .= sol
    end

### do stuff with pre_alloc_sols which returns some loss
 
     return loss
end

opt_func = Optimization.OptimizationFunction((x, p, X, Y) -> loss(x, X, Y, num_shots), ...)
opt_prob = Optimization.OptimizationProblem(opt_func, ...)

res = Optimization.solve(opt_prob, ...)

I’ve omitted a lot of detail so obviously the code won’t work and isn’t a MWE, but I basically don’t know what typeToInstantiate. I could probably do Any but that feels like it’d be inefficient. Any tips appreciated.

2

You’ll need to make use of PreallocationTools.jl:

2 Likes