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

**URL:** https://discourse.julialang.org/t/pre-allocating-arrays-for-parallel-multiple-shooting-neural-ode-training-in-flux/97867
**Category:** Machine Learning
**Tags:** flux, optimization, differentialequation
**Created:** [April 24, 2023, 4:43pm UTC](https://discourse.julialang.org/t/pre-allocating-arrays-for-parallel-multiple-shooting-neural-ode-training-in-flux/97867 "2023-04-24T16:43:22Z")
**Posts on this page:** 2
**Page:** 1

<div class="post-metadata">

### Author: ![linkz](https://avatars.discourse-cdn.com/v4/letter/l/da6949/32.png) [@linkz](https://discourse.julialang.org/u/linkz)
#### Post date: [April 24, 2023, 4:43pm UTC](https://discourse.julialang.org/t/pre-allocating-arrays-for-parallel-multiple-shooting-neural-ode-training-in-flux/97867/1 "2023-04-24T16:43:22Z")

</div>

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:

```julia
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.

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [April 24, 2023, 5:30pm UTC](https://discourse.julialang.org/t/pre-allocating-arrays-for-parallel-multiple-shooting-neural-ode-training-in-flux/97867/2 "2023-04-24T17:30:39Z")

</div>

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

> **[GitHub - SciML/PreallocationTools.jl: Tools for building non-allocating...](https://github.com/SciML/PreallocationTools.jl)**
>
> Tools for building non-allocating pre-cached functions in Julia, allowing for GC-free usage of automatic differentiation in complex codes - GitHub - SciML/PreallocationTools.jl: Tools for building ...