在 Julia 中定义自动广播的分段函数

Question

我有一个类似于下面

f

形式的函数。

function f(x, y)
    if x <= y
        return x + y
    else
        return x - y
    end
end

我想修改它，以便它可以接受混合标量和向量，例如

f(1, 2) = 3

和

f([1, 4], 2) = [3, 2]

。实现这一目标并保持其良好的 Latexify.jl 输出的最简单方法是什么？

尝试：

• 在调用站点进行广播是可行的，但我宁愿在定义中处理它。

• @. if x <= y

  会产生错误，因为

  if([1,0])

  不起作用。
• 我想过使用数组理解，但那永远不会返回标量。

Answer 1

您可以定义多个重载来处理接收数组：

f_broadcast(x, y) = ifelse.(x .<= y, x .+ y, x .- y)
f(x::AbstractArray, y) = f_broadcast(x, y)
f(x, y::AbstractArray) = f_broadcast(x, y)
f(x::AbstractArray, y::AbstractArray) = f_broadcast(x, y)
f(x, y) = ifelse(x <= y, x + y, x - y)

@show f(1, 2) f([1, 4], 2) f(2, [1, 4]) f([1, 2], [3, 4])

f(1, 2) = 3
f([1, 4], 2) = [3, 2]
f(2, [1, 4]) = [1, 6]
f([1, 2], [3, 4]) = [4, 6]