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?
# 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
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:
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