Handle Base.hypot with more than 2 arguments

Author: moble

src/rulesets/Base/base.jl

@@ -119,6 +119,40 @@ function rrule(::typeof(hypot), z::Complex)
     return (Ω, hypot_pullback)
 end
 
+# Note that `hypot` with two arguments has rules in `fastmath_able.jl`
+
+function frule(
+    (_, Δx, Δy, Δz, Δxs...),
+    ::typeof(hypot),
+    x::T,
+    y::T,
+    z::T,
+    xs::Vararg{T, N}
+) where {T<:Union{Real,Complex}, N}
+    Ω = hypot(x, y, z, xs...)
+    n = ifelse(iszero(Ω), one(Ω), Ω)
+    ∂Ωxyz = realdot(x, Δx) + realdot(y, Δy) + realdot(z, Δz)
+    ∂Ωxs = sum(realdot(xi, Δxi) for (xi, Δxi) in zip(xs, Δxs); init=zero(∂Ωxyz))
+    ∂Ω = (∂Ωxyz + ∂Ωxs) / n
+    return Ω, ∂Ω
+end
+
+function rrule(
+    ::typeof(hypot),
+    x::T,
+    y::T,
+    z::T,
+    xs::Vararg{T, N}
+) where {T<:Union{Real,Complex}, N}
+    Ω = hypot(x, y, z, xs...)
+    function hypot_pullback(ΔΩ)
+        c = real(ΔΩ) / ifelse(iszero(Ω), one(Ω), Ω)
+        return (NoTangent(), c * x, c * y, c * z, (c * xi for xi in xs)...)
+    end
+    return (Ω, hypot_pullback)
+end
+
+
 @scalar_rule fma(x, y, z) (y, x, true)
 @scalar_rule muladd(x, y, z) (y, x, true)
 @scalar_rule muladd(x::Union{Number, ZeroTangent}, y::Union{Number, ZeroTangent}, z::Union{Number, ZeroTangent}) (y, x, true)

test/rulesets/Base/base.jl

@@ -263,4 +263,15 @@ end
         test_rrule(merge, (; a=1.0), (; b=2.0))
         test_rrule(merge, (; a=1.0), (; a=2.0))
     end
+
+    @testset "hypot(x, y, z, xs...)" begin
+        test_frule(hypot, 1.0, 2.0, 3.0)
+        test_rrule(hypot, 1.0, 2.0, 3.0)
+        test_frule(hypot, 1.0, 2.0, 3.0, 4.0)
+        test_rrule(hypot, 1.0, 2.0, 3.0, 4.0)
+        test_frule(hypot, 1.0+5.0im, 2.0+6.0im, 3.0+7.0im)
+        test_rrule(hypot, 1.0+5.0im, 2.0+6.0im, 3.0+7.0im)
+        test_frule(hypot, 1.0+5.0im, 2.0+6.0im, 3.0+7.0im, 4.0+8.0im)
+        test_rrule(hypot, 1.0+5.0im, 2.0+6.0im, 3.0+7.0im, 4.0+8.0im)
+    end
 end