Add rules for hypot (#2626)

* Add rules for hypot

* Don't bother storing hypot value in tape

* Format with runic

* Deal with complex values

* Reduce duplication by rearranging if statements

* Avoid NaNs at the origin

* Change hypot rules to apply with at least 3 arguments

---------

Co-authored-by: Mosè Giordano <[email protected]>

Author: moble

src/internal_rules.jl

@@ -1647,3 +1647,107 @@ function EnzymeRules.reverse(
     dst::Annotation{<:AbstractArray})
     return (nothing, nothing)
 end
+
+
+_hypotforward(x::Const) = zero(x.val)
+_hypotforward(x) = real(conj(x.val) * x.dval)
+_hypotforward(x::Const, i) = zero(x.val)
+_hypotforward(x, i) = real(conj(x.val) * x.dval[i])
+
+function EnzymeRules.forward(
+        config::EnzymeRules.FwdConfig,
+        func::Const{typeof(Base.hypot)},
+        RT,
+        x::Annotation,
+        y::Annotation,
+        z::Annotation,
+        xs::Vararg{Annotation, N}
+    ) where {N}
+    if EnzymeRules.needs_shadow(config)
+        h = func.val(x.val, y.val, z.val, map(x -> x.val, xs)...)
+        n = iszero(h) ? one(h) : h
+        if EnzymeRules.width(config) == 1
+            dh = (
+                _hypotforward(x) +
+                    _hypotforward(y) +
+                    _hypotforward(z) +
+                    sum(_hypotforward, xs, init = zero(real(x.val)))
+            ) / n
+            if EnzymeRules.needs_primal(config)
+                return Duplicated(h, dh)
+            else
+                return dh
+            end
+        else
+            dh = ntuple(
+                i -> (
+                    _hypotforward(x, i) +
+                        _hypotforward(y, i) +
+                        _hypotforward(z, i) +
+                        sum(x -> _hypotforward(x, i), xs; init = zero(real(x.val)))
+                ) / n,
+                Val(EnzymeRules.width(config)),
+            )
+            if EnzymeRules.needs_primal(config)
+                return BatchDuplicated(h, dh)
+            else
+                return dh
+            end
+        end
+    elseif EnzymeRules.needs_primal(config)
+        return func.val(x.val, y.val, z.val, map(x -> x.val, xs)...)
+    else
+        return nothing
+    end
+end
+
+_hypotreverse(x::Const, ::Val{W}, dret::Const, h) where {W} = nothing
+_hypotreverse(x::Const, ::Val{W}, dret, h) where {W} = nothing
+function _hypotreverse(x, w::Val{W}, dret::Const, h) where {W}
+    if W == 1
+        return zero(x.val)
+    else
+        return ntuple(Returns(zero(x.val)), w)
+    end
+end
+function _hypotreverse(x, w::Val{W}, dret, h) where {W}
+    if W == 1
+        return x.val * dret.val / h
+    else
+        return ntuple(i -> x.val * dret.val[i] / h, w)
+    end
+end
+
+function EnzymeRules.augmented_primal(
+        config::EnzymeRules.RevConfig,
+        func::Const{typeof(Base.hypot)},
+        ::Type,
+        x::Annotation,
+        y::Annotation,
+        z::Annotation,
+        xs::Vararg{Annotation, N}
+    ) where {N}
+    h = hypot(x.val, y.val, z.val, map(x -> x.val, xs)...)
+    primal = needs_primal(config) ? h : nothing
+    return EnzymeRules.AugmentedReturn(primal, nothing, nothing)
+end
+
+function EnzymeRules.reverse(
+        config::EnzymeRules.RevConfig,
+        func::Const{typeof(Base.hypot)},
+        dret,
+        tape,
+        x::Annotation,
+        y::Annotation,
+        z::Annotation,
+        xs::Vararg{Annotation, N}
+    ) where {N}
+    h = hypot(x.val, y.val, z.val, map(x -> x.val, xs)...)
+    n = iszero(h) ? one(h) : h
+    w = Val(EnzymeRules.width(config))
+    dx = _hypotreverse(x, w, dret, n)
+    dy = _hypotreverse(y, w, dret, n)
+    dz = _hypotreverse(z, w, dret, n)
+    dxs = map(x -> _hypotreverse(x, w, dret, n), xs)
+    return (dx, dy, dz, dxs...)
+end

test/rules/internal_rules.jl

@@ -830,4 +830,44 @@ end
 
 end
 
+
+@testset "hypot rules" begin
+    @testset "forward" begin
+        @testset for RT in (Const, DuplicatedNoNeed, Duplicated),
+                Tx in (Const, Duplicated),
+                Ty in (Const, Duplicated),
+                Tz in (Const, Duplicated),
+                Txs in (Const, Duplicated)
+
+            x, y, z, xs = 2.0, 3.0, 5.0, 17.0
+            test_forward(hypot, RT, (x, Tx), (y, Ty))
+            test_forward(hypot, RT, (x, Tx), (y, Ty), (z, Tz))
+            test_forward(hypot, RT, (x, Tx), (y, Ty), (z, Tz), (xs, Txs))
+
+            x, y, z, xs = 2.0 + 7.0im, 3.0 + 11.0im, 5.0 + 13.0im, 17.0 + 19.0im
+            test_forward(hypot, RT, (x, Tx), (y, Ty))
+            test_forward(hypot, RT, (x, Tx), (y, Ty), (z, Tz))
+            test_forward(hypot, RT, (x, Tx), (y, Ty), (z, Tz), (xs, Txs))
+        end
+    end
+    @testset "reverse" begin
+        @testset for RT in (Active,),
+                Tx in (Const, Active),
+                Ty in (Const, Active),
+                Tz in (Const, Active),
+                Txs in (Const, Active)
+
+            x, y, z, xs = 2.0, 3.0, 5.0, 17.0
+            test_reverse(hypot, RT, (x, Tx), (y, Ty))
+            test_reverse(hypot, RT, (x, Tx), (y, Ty), (z, Tz))
+            test_reverse(hypot, RT, (x, Tx), (y, Ty), (z, Tz), (xs, Txs))
+
+            x, y, z, xs = 2.0 + 7.0im, 3.0 + 11.0im, 5.0 + 13.0im, 17.0 + 19.0im
+            test_reverse(hypot, RT, (x, Tx), (y, Ty))
+            test_reverse(hypot, RT, (x, Tx), (y, Ty), (z, Tz))
+            test_reverse(hypot, RT, (x, Tx), (y, Ty), (z, Tz), (xs, Txs))
+        end
+    end
+end
+
 end # InternalRules