Search code examples
juliavectorization

Vectorizing a multiple output function in Julia

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

Solution

  • 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])