Style question: replacing multidimensional collect with broadcast

What’s the recommended way of writing something like

[f(i, j) for i in 1:I, j in 1:J]

using broadcasting? Eg

f.(1:I, (1:J)')
f.(1:I, permutedims(1:J))

?

This is purely a style question. Transpose feels icky (relies on 1'==1), permutedims is long.

1 Like

I personally like the icky version more :wink: (comprehension even better, despite the performance issue)

How about this outer() function? Outer broadcasting? - #6 by rdeits

Or we could turn that into a macro so that the result would fuse (I think) with any outer broadcasting operations.

2 Likes

I am not sure there is a performance issue; differences are minor and the order depends on the size.

Thanks for linking that topic, I see that this is an open question. Yes, somehow having a shorthand that combines well with broadcasting would be nice.

How about a syntax f.↑(1:I, 1:J) so that you can do, e.g.,

julia> tuple.↑([1,2,3],[4,5])
3×2 Array{Tuple{Int64,Int64},2}:
 (1, 4)  (1, 5)
 (2, 4)  (2, 5)
 (3, 4)  (3, 5)

Here is an implementation based on @rdeits’s outer:

@generated function outer_broadcasted(f, args::Tuple{Vararg{Any, N}}) where {N}
    expr = Expr(:call, Broadcast.broadcasted, :f)  # instead of broadcast
    for i in 1:N
        if i == 1
            push!(expr.args, :(args[1]))
        else
            push!(expr.args, Expr(:call, :reshape, :(args[$i]), Expr(:tuple, [1 for _ in 1:i - 1]..., :(:))))
        end
    end
    expr
end

↑(f, args::Tuple) = Broadcast.materialize(f.↑(args))
Broadcast.broadcasted(::typeof(↑), f, args) = outer_broadcasted(f, args)

It’ll be fused with other dotted functions:

julia> Meta.@lower f.↑(x, y) .+ z
:($(Expr(:thunk, CodeInfo(
1 ─ %1 = (Core.tuple)(x, y)
│   %2 = (Base.broadcasted)(↑, f, %1)
│   %3 = (Base.broadcasted)(+, %2, z)
│   %4 = (Base.materialize)(%3)
└──      return %4
))))

A downside is that f.↑(x, y) .^ z does not work as you’d expect.

1 Like

How about making vec(x, d) create a vector along dimension d? (Not type-stable in the general case, but d would typically be a literal.) Then one would write f.(1:I, vec(1:J, 2)), which is slightly shorter than premutedims and also generalizes better to higher dimension.

For an even shorter alternative, map some Unicode operator to vec so that one can write f.(1:I, 1:J⇾2) or similar.

2 Likes

And then make

vec(n::Int, dim::Int) = vec(Base.OneTo(n), dim)

I like this very much. Maybe I will experiment with this as a mini-package, also solving the n x 1 array creation problem.