Vectorizing a multiple output function in Julia

ValientProcess 0
Viewed 81

I'm trying to vectorize a multiple output function in Julia, for example this function:

function z(a,b)
    x_1 = a*b
    x_2 = a + b
    x_3 = a^b
    return x_1, x_2, x_3
end

With the input:

begin
B = [1, 4]
A = 2
end

I want the output to be x_1 = (2,8), x_2 = (3,6), x_3 = (2,16). However, the default vectorization z.(A,B) returns this:

[(2,3,2),(8,6,16)]

I wonder if there is a quick and efficient way to make the function return the values in the shape I need, rather than manipulating the output after calling the function

vectorization julia
2 Answers

You can apply the vectorization inside the function definition.

julia> function z(a,b)
           x_1 = a .* b
           x_2 = a .+ b
           x_3 = a .^ b
           return x_1, x_2, x_3
       end
z (generic function with 1 method)

julia> A = 2
2

julia> B = [1, 4]
2-element Vector{Int64}:
 1
 4

julia> z(A, B)
([2, 8], [3, 6], [2, 16])
BallpointBen0

This can be done quite nicely with StructArrays.jl:

julia> using StructArrays

julia> function unzip_broadcast(f, args...)
         # first use Julia's own lazy broadcasting:
         bc = Broadcast.instantiate(Broadcast.broadcasted(f, args...))
         # but instead of one array of tuples, make a tuple of arrays:
         StructArrays.components(StructArray(bc))
       end;

julia> unzip_broadcast(z, [1,4], 2)
([2, 8], [3, 6], [1, 16])

This is about as fast as BallpointBen's answer which re-writes the function to broadcast inside, on this particular example. But it can be used for functions which you can't change.

mcabbott0
Related