Providing hints to broadcast?

New to Julia
#1

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”???

#2

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

1 Like
#3

Lift b and c into that third dimension:

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

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

#5

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
#6

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

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

#8

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.

#9

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