Handle Base.hypot with more than 2 arguments (#824)

* Handle Base.hypot with more than 2 arguments

* Bump version

* Apply suggestions from code review

Co-authored-by: David Müller-Widmann <[email protected]>

* Test mixed Real/Complex inputs to hypot

---------

Co-authored-by: David Müller-Widmann <[email protected]>

Author: moble

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.72.5"
+version = "1.72.6"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/rulesets/Base/base.jl

@@ -119,6 +119,38 @@ 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::Union{Real,Complex},
+    y::Union{Real,Complex},
+    z::Union{Real,Complex},
+    xs::Union{Real,Complex}...,
+)
+    Ω = hypot(x, y, z, xs...)
+    n = ifelse(iszero(Ω), oneunit(Ω), Ω)
+    ∂Ω = sum(map(realdot, (x, y, z, xs...), (Δx, Δy, Δz, Δxs...))) / n
+    return Ω, ∂Ω
+end
+
+function rrule(
+    ::typeof(hypot),
+    x::Union{Real,Complex},
+    y::Union{Real,Complex},
+    z::Union{Real,Complex},
+    xs::Union{Real,Complex}...,
+)
+    Ω = hypot(x, y, z, xs...)
+    function hypot_pullback(ΔΩ)
+        c = real(ΔΩ) / ifelse(iszero(Ω), one(Ω), Ω)
+        return (NoTangent(), c * x, c * y, c * z, map(xi -> c * xi, 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,14 @@ 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
+        for n in (3, 4)
+            for sig in Iterators.product(ntuple(_->(Float64, ComplexF64), n)...)
+                args = randn.(sig)
+                test_frule(hypot, args...)
+                test_rrule(hypot, args...)
+            end
+        end
+    end
 end