在特定参数上矢量化函数

Question

假设我有一个功能

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

我想仅对第二个参数进行矢量化,以便我有类似的东西

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

我知道有宏@ vectorize_1arg或@ vectorize_2arg.但是,这些宏将向量化所有参数.

问题:如何方便地对特定参数的函数进行矢量化？我必须像上面的例子一样硬编码吗？

Answer 1

如果您希望扩展只需要第二个参数向量化的函数，那么应该这样做：

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

像这样使用：

@vectorize_on_2nd Int64 myfunc

那应该给你一个myfunc{N}(::Any, ::AbstractArray{Int64,N})方法。