Vectorize a function on a specific argument

Question

Suppose I have a function

myfunc(a, x::Int64) = a * x

I want to vectorize the second argument only, so that I have something like

myfunc{N}(a, x::Array{Int64, N}) = map(x -> myfunc(a, x), x)

I know there are macro @vectorize_1arg or @vectorize_2arg. However, those macros will vectorize all arguments.

Question: How to vectorize the function on a specific argument conveniently? Do I have to hard code like the above example?

Answer 1

If you're looking to extend functions where you want only the second arg vectorized, this should do it:

macro vectorize_on_2nd(S, f)
    S = esc(S); f = esc(f); N = esc(:N)
    quote
        ($f){$N}(a, x::AbstractArray{$S, $N}) =
            reshape([($f)(a, x[i]) for i in eachindex(x)], size(x))
    end
end

Used like this:

@vectorize_on_2nd Int64 myfunc

That should give you a myfunc{N}(::Any, ::AbstractArray{Int64,N}) method.

Answer 2

Most of the time, this works

myfunc.([a],x)

as it will be vectorized over both arguments but [a] has only one entry.