# Unstack-stack-unstack-stack

**URL:** <https://discourse.julialang.org/t/unstack-stack-unstack-stack/43591>\
**Category:** Machine Learning\
**Tags:** question\
**Created:** [July 23, 2020, 9:57pm UTC](https://discourse.julialang.org/t/unstack-stack-unstack-stack/43591 "2020-07-23T21:57:55Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![victorialena](https://avatars.discourse-cdn.com/v4/letter/v/9de0a6/32.png) [@victorialena](https://discourse.julialang.org/u/victorialena)\
**Post date:** [July 23, 2020, 9:57pm UTC](https://discourse.julialang.org/t/unstack-stack-unstack-stack/43591/1 "2020-07-23T21:57:55Z")

</div>

any ideas on how to make this more efficient?  
I don’t know how to broadcast the functions more efficiently than a for loop and matrix multiplications in 3D didn’t really work for me. Can someone help?

```julia
# input has size [Ni, B]
# solutions is a Array{Function} of length Np
b = model_one(input) #output is (Na, B)
w = model_two(input) #output is (Np+1, B)
B, Np = size(input, 2), size(w,1) - 1
    
qs = Zygote.ignore(() -> [hcat([s(input[:,i]) for s in solutions]...) for i=1:B]) #output is Bx(Na, Np)
Flux.stack(qs .* Flux.unstack(w[1:Np, :], 2), 2) .+ w[Np+1:Np+1, :] .* b

```

---

<div class="post-metadata">

**Author:** ![SteffenPL](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/steffenpl/32/206270_2.png) [@SteffenPL](https://discourse.julialang.org/u/SteffenPL)\
**Post date:** [July 25, 2020, 3:45pm UTC](https://discourse.julialang.org/t/unstack-stack-unstack-stack/43591/2 "2020-07-25T15:45:45Z")

</div>

I’m not sure if I understand exactly the context.

- From the code I guessed that `s(input[:,i])` has dimension `Na`.
- Do you want to get more speed or a better readable code?

My recommendation is to use a plain for-loop for your task.  
This makes it easier to understand and probably even faster. In particular, it is not needed to allocate an array of size `(Na, Np, B)` for your calculation.

But, I’m not sure if the Zygote context allows this for-loop. If array indexing is not allowed for your context, we can try to find another solution.

Here is a minimal example:

```julia
Ni = 3
Na = 5
Np = 7
B = 11

input = rand(Ni,B)
solutions = [x -> x[1] .* rand(Na) for _ in 1:Np]

b = rand(Na,B)
w = rand(Np+1,B)

res = zeros(Na,B)
for j = 1:B
    global res # just for the minimal example to work in global scope
    for i = 1:Np
        res[:,j] += w[i,j] .* solutions[i](input[:,j])
    end
    res[:,j] += w[end,j] .* b[:,j] # this operation could also be done outside of the loop, like in your code
end

```