vectorization - 对特定参数的函数进行向量化

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

如果您想在只需要向量化第二个 arg 的地方扩展函数，则应该这样做:

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})方法。