在特定参数上矢量化函数

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

Answer 2

大多数情况下，这种方法有效

myfunc.([a],x)

因为它将在两个参数上向量化，但是[a]只有一个条目。