Macro that vectorizes function with two args of different types

[Boris Kheyfets]

I wanted to apply vectrorize_2arg macro to a function with two args of different type:

function mysquare(n::Int64, x::Float64)
    return n * x^2
end

println(mysquare(2, 2.))

@vectorize_2arg Int64 Float64 mysquare # ⇒ wrong number of args

println(mysquare(2, [2., 3., 4.]))

I modified the code of vectorize_2arg macro and it seems to be working: (at least with tests I tried)

macro vectorize_2arg_full(S1,S2,f)
    S1 = esc(S1); S2 = esc(S2); f = esc(f); T1 = esc(:T1); T2 = esc(:T2)
    quote
        ($f){$T1<:$S1, $T2<:$S2}(x::($T1), y::AbstractArray{$T2}) =
            reshape([ ($f)(x, y[i]) for i=1:length(y) ], size(y))
        ($f){$T1<:$S1, $T2<:$S2}(x::AbstractArray{$T1}, y::($T2)) =
            reshape([ ($f)(x[i], y) for i=1:length(x) ], size(x))

        function ($f){$T1<:$S1, $T2<:$S2}(x::AbstractArray{$T1}, y::AbstractArray{$T2})
            shp = promote_shape(size(x),size(y))
            reshape([ ($f)(x[i], y[i]) for i=1:length(x) ], shp)
        end
    end
end

function mysquare(n::Int64, x::Float64)
    return n * x^2
end

println(mysquare(2, 2.))

@vectorize_2arg_full Int64 Float64 mysquare

println(mysquare(2, [2., 3., 4.]))

Can someone code-review it?

[Stefan Karpinski]

Why not call the macro with Number as the type argument?

[Boris Kheyfets]

Why not call the macro with Number as the type argument?

My bad: I didn't figure it out before.

Anyway, it seems like my version is slightly faster:

#!/usr/bin/env julia

macro vectorize_2arg_full(S1,S2,f)
    S1 = esc(S1); S2 = esc(S2); f = esc(f); T1 = esc(:T1); T2 = esc(:T2)
    quote
        ($f){$T1<:$S1, $T2<:$S2}(x::($T1), y::AbstractArray{$T2}) =
            reshape([ ($f)(x, y[i]) for i=1:length(y) ], size(y))
        ($f){$T1<:$S1, $T2<:$S2}(x::AbstractArray{$T1}, y::($T2)) =
            reshape([ ($f)(x[i], y) for i=1:length(x) ], size(x))

        function ($f){$T1<:$S1, $T2<:$S2}(x::AbstractArray{$T1}, y::AbstractArray{$T2})
            shp = promote_shape(size(x),size(y))
            reshape([ ($f)(x[i], y[i]) for i=1:length(x) ], shp)
        end
    end
end

function mysquare_1(n::Int64, x::Float64)

    return n * x^2
end

@vectorize_2arg Number mysquare_1

function mysquare_2(n::Int64, x::Float64)

    return n * x^2
end

@vectorize_2arg_full Int64 Float64 mysquare_2

println( @time mysquare_1(2, [2., 3., 4.]) )
println( @time mysquare_2(2, [2., 3., 4.]) )

Gives:

elapsed time: 0.009493773 seconds (343832 bytes allocated)
[8.0,18.0,32.0]
elapsed time: 0.006558491 seconds (232872 bytes allocated)
[8.0,18.0,32.0]

[Stefan Karpinski]

You're timing the JIT compilation. The second function shares code with the first one, so compiling it takes less time. Run it again in the same session and you'll see that this computation is too fast to meaningfully time:

julia> println( @time mysquare_1(2, [2., 3., 4.]) )

elapsed time: 7.226e-6 seconds (416 bytes allocated)

[8.0,18.0,32.0]

julia> println( @time mysquare_2(2, [2., 3., 4.]) )

elapsed time: 9.414e-6 seconds (416 bytes allocated)

[8.0,18.0,32.0]

The difference here is noise and the time is close to the computer's time resolution.

[Boris Kheyfets]

You right. My mistake.