I often run into this situation and have yet to find a solution other than writing a for-loop.
Is there a way to vectorize the MWE below using features in Base?

Suppose I have a function accepting a 2d array and two integers that returns a transformed array:

``````function foo(a::Array{Float64,2}, b::Int, c::Int)
#some transform
return newa
end
``````

I want to be able to generalize the usage of `foo()` with a broadcast:

``````a = ones(5,5)
b = [1,2,3]
c = [1,2,3]
result = zeros(5,5,3)
result .= foo.(a, b, c)
``````

In this case I would intuit that `a` gets expanded on its 3rd dimension. However, broadcast can’t handle that. The closest one-liner solution I can come up with is:

``````result = foo.([a],b,c)
``````

…but this returns an Array of Arrays which obviously isn’t great either.

I had thought the fairly recent addition of `eachslice()` might help but it doesn’t so much in this situation.

Any ideas? Is it possible to provide a clue to broadcast that indicates how the first Array should be expanded? Or is it “just write a for loop”???

One way, probably not so efficient, is to add a singleton dimension to hint the broadcast:

1 Like

Lift `b` and `c` into that third dimension:

``````result .= foo.(a, reshape(b, 1, 1, :), reshape(c, 1, 1, :))
``````
2 Likes

Both replies involve using `reshape` to “add” the missing dimension. However, this (surprisingly) also doesn’t work.

Oh, I missed the signature of `foo(a::Array{Float64,2}, b::Int, c::Int)`. My answer assumes that `foo` is truly a scalar operation and operates element wise across all three arrays. You want to operate element wise across `b` and `c`, but treat `a` as an array, and concatenate the answers into a third dimension. That’s simply not what broadcast does.

2 Likes

Yeah, I assumed it was pushing things into odd territory. But seeing as I often pleasantly find clever (often undocumented) tricks, worth an ask. I will likely just write an extra dispatch that envelops the original function in a for loop. Thanks!

1 Like

Does Ref work?
`foo.(Ref(a),b,c)`

It’s the concatenating part that’s the challenge to express with broadcast alone.

Using Ref is the same as doing `[a]`, which was tried in the original post and returns an array of arrays. You can always concatenate them together afterwards, though.

A nice way to combine slices is the package JuliennedArrays:

``````julia> foo(a::Array{Float64,2}, b::Int, c::Int) = a .+ (b+c); # for example

julia> using JuliennedArrays

julia> collect(Align(foo.(Ref(a),b,c), 2,3))
3×5×5 Array{Float64,3}:
[:, :, 1] =
3.0  3.0  3.0  3.0  3.0
5.0  5.0  5.0  5.0  5.0
7.0  7.0  7.0  7.0  7.0

[:, :, 2] =
3.0  3.0  3.0  3.0  3.0

``````

I think you can also write `reshape(reduce(hcat, foo.(Ref(a),b,c) ), (5,5,3))` but it’s possible I messed up the order of sizes here.

1 Like