Add derivative rules for norm and fix hypot at zero

Project.toml

@@ -15,10 +15,12 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [weakdeps]
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [extensions]
 ForwardDiffStaticArraysExt = "StaticArrays"
+SparseArraysExt = "SparseArrays"
 
 [compat]
 Calculus = "0.5"
@@ -31,6 +33,7 @@ JET = "0.9, 0.12, 0.13"
 LogExpFunctions = "0.3, 1"
 NaNMath = "1"
 Preferences = "1"
+SparseArrays = "1.10"
 SpecialFunctions = "1, 2"
 StaticArrays = "1.5"
 julia = "1.10"

docs/src/user/limitations.md

@@ -7,7 +7,7 @@ function being differentiated):
 
 - **The target function can only be composed of generic Julia functions.** ForwardDiff cannot propagate derivative information through non-Julia code. Thus, your function may not work if it makes calls to external, non-Julia programs, e.g. uses explicit BLAS calls instead of `Ax_mul_Bx`-style functions.
 
-- **The target function must be a composition of differentiable functions.** ForwardDiff can have issues to compute derivatives of functions, which are composed of at least one function, which is not differentiable in the point the derivative should be evaluated, even if the target function itself is differentiable. A simple example is `f(x) = norm(x)^2`, where `ForwardDiff.gradient(f, zeros(2))` returns a vector of `NaN`s since the Euclidean norm is not differentiable in zero. A possible solution to this issue is to, e.g., define `f(x) = sum(abs2, x)` instead. In situations, where rewriting the target function only as a composition of differentiable functions is more complicated (e.g. `f(x) = (1 + norm(x))exp(-norm(x))`)), one would need to define a custom derivative rule (see [this comment](https://github.com/JuliaDiff/ForwardDiff.jl/issues/303#issuecomment-2977990425)).
+- **The target function must be a composition of differentiable functions.** ForwardDiff can have issues to compute derivatives of functions, which are composed of at least one function, which is not differentiable in the point the derivative should be evaluated, even if the target function itself is differentiable. A simple example is `f(x) = norm(x)^2`: Since the Euclidean norm is not differentiable in zero, ForwardDiff uses the subgradient zero of `norm` at zero, which is then propagated through all subsequent functions. While `ForwardDiff.gradient(f, zeros(2))` happens to be correct, `ForwardDiff.hessian(f, zeros(2))` returns a matrix of zeros instead of `2I`. A possible solution to this issue is to, e.g., define `f(x) = sum(abs2, x)` instead. In situations, where rewriting the target function only as a composition of differentiable functions is more complicated (e.g. `f(x) = (1 + norm(x))exp(-norm(x))`)), one would need to define a custom derivative rule (see [this comment](https://github.com/JuliaDiff/ForwardDiff.jl/issues/303#issuecomment-2977990425)).
 
 - **The target function must be unary (i.e., only accept a single argument).** [`ForwardDiff.jacobian`](@ref) is an exception to this rule.
 

ext/ForwardDiffStaticArraysExt.jl

@@ -42,6 +42,13 @@ function LinearAlgebra.eigen(A::Hermitian{Complex{Dual{T,V,N}}, <:StaticArrays.S
     return ForwardDiff._eigen_hermitian(A)
 end
 
+const StaticDualArrayRealComplex{T,V,N} = Union{StaticArray{<:Tuple,Dual{T,V,N}},StaticArray{<:Tuple,Complex{Dual{T,V,N}}}}
+LinearAlgebra.norm(x::StaticDualArrayRealComplex, p::Real=2) = ForwardDiff._norm_dual_x(x, p)
+LinearAlgebra.norm(x::StaticArray{<:Tuple,<:Number}, p::Dual) = ForwardDiff._norm_dual_p(x, p)
+LinearAlgebra.norm(x::StaticDualArrayRealComplex, p::Dual) = ForwardDiff._norm_dual_xp(x, p)
+# Faster than broadcasting
+ForwardDiff._sum_map(f::F, x::StaticArray, vx::StaticArray) where {F} = ForwardDiff._sum_map(f, Tuple(x), Tuple(vx))
+
 # For `MMatrix` we can use the in-place method
 ForwardDiff._lyap_div_zero_diag!!(A::StaticArrays.MMatrix, λ::AbstractVector) = ForwardDiff._lyap_div_zero_diag!(A, λ)
 

ext/SparseArraysExt.jl

@@ -0,0 +1,22 @@
+module SparseArraysExt
+
+using ForwardDiff: Dual
+using LinearAlgebra: LinearAlgebra, norm
+using SparseArrays: SparseArrays, AbstractSparseMatrixCSC, SparseVectorUnion
+
+# Resolve method ambiguities
+for A in (:AbstractSparseMatrixCSC, :SparseVectorUnion)
+    @eval begin
+        function LinearAlgebra.norm(x::Union{$A{Dual{T,V,N}},$A{Complex{Dual{T,V,N}}}}, p::Real=2) where {T,V,N}
+            return invoke(norm, Tuple{$A,Real}, x, p)
+        end
+        function LinearAlgebra.norm(x::$A{<:Number}, p::Dual)
+            return invoke(norm, Tuple{$A,Real}, x, p)
+        end
+        function LinearAlgebra.norm(x::Union{$A{Dual{T,V,N}},$A{Complex{Dual{T,V,N}}}}, p::Dual) where {T,V,N}
+            return invoke(norm, Tuple{$A,Real}, x, p)
+        end
+    end
+end
+
+end # module

src/dual.jl

@@ -474,7 +474,7 @@ Base.float(d::Dual) = convert(float(typeof(d)), d)
 ###################################
 
 for (M, f, arity) in DiffRules.diffrules(filter_modules = nothing)
-    if (M, f) in ((:Base, :^), (:NaNMath, :pow), (:Base, :/), (:Base, :+), (:Base, :-), (:Base, :sin), (:Base, :cos))
+    if (M, f) in ((:Base, :^), (:NaNMath, :pow), (:Base, :/), (:Base, :+), (:Base, :-), (:Base, :sin), (:Base, :cos), (:Base, :hypot))
         continue  # Skip methods which we define elsewhere.
     elseif !(isdefined(@__MODULE__, M) && isdefined(getfield(@__MODULE__, M), f))
         continue  # Skip rules for methods not defined in the current scope
@@ -599,26 +599,186 @@ end
 # hypot #
 #-------#
 
-@inline function calc_hypot(x, y, z, ::Type{T}) where T
-    vx = value(x)
-    vy = value(y)
-    vz = value(z)
-    h = hypot(vx, vy, vz)
-    p = (vx / h) * partials(x) + (vy / h) * partials(y) + (vz / h) * partials(z)
-    return Dual{T}(h, p)
-end
+# `hypot` and `norm` are not differentiable at zero; we use the subgradient zero there
+_iszero_primal(x::Dual) = _iszero_primal(x.value)
+_iszero_primal(x) = iszero(x)
+_div_or_zero(v, h) = _iszero_primal(h) ? zero(v / oneunit(h)) : v / h
+
+const DualNTupleRealComplex{T,V,N} = Union{
+    Tuple{Dual{T,V,N},Vararg{Dual{T,V,N}}},Tuple{Complex{Dual{T,V,N}},Vararg{Complex{Dual{T,V,N}}}}
+}
+
+# Arguments that are not Duals of the result tag are constants and are passed through unchanged
+@define_binary_dual_op(
+    Base.hypot,
+    begin
+        h = hypot(x.value, y.value)
+        dual_definition_retval(Val{Txy}(), h, _div_or_zero(x.value, h), x.partials, _div_or_zero(y.value, h), y.partials)
+    end,
+    begin
+        h = hypot(x.value, y)
+        dual_definition_retval(Val{Tx}(), h, _div_or_zero(x.value, h), x.partials)
+    end,
+    begin
+        h = hypot(x, y.value)
+        dual_definition_retval(Val{Ty}(), h, _div_or_zero(y.value, h), y.partials)
+    end,
+)
 
 @define_ternary_dual_op(
     Base.hypot,
-    calc_hypot(x, y, z, Txyz),
-    calc_hypot(x, y, z, Txy),
-    calc_hypot(x, y, z, Txz),
-    calc_hypot(x, y, z, Tyz),
-    calc_hypot(x, y, z, Tx),
-    calc_hypot(x, y, z, Ty),
-    calc_hypot(x, y, z, Tz),
+    begin
+        h = hypot(x.value, y.value, z.value)
+        p = _div_or_zero(x.value, h) * x.partials + _div_or_zero(y.value, h) * y.partials + _div_or_zero(z.value, h) * z.partials
+        Dual{Txyz}(h, p)
+    end,
+    begin
+        h = hypot(x.value, y.value, z)
+        dual_definition_retval(Val{Txy}(), h, _div_or_zero(x.value, h), x.partials, _div_or_zero(y.value, h), y.partials)
+    end,
+    begin
+        h = hypot(x.value, y, z.value)
+        dual_definition_retval(Val{Txz}(), h, _div_or_zero(x.value, h), x.partials, _div_or_zero(z.value, h), z.partials)
+    end,
+    begin
+        h = hypot(x, y.value, z.value)
+        dual_definition_retval(Val{Tyz}(), h, _div_or_zero(y.value, h), y.partials, _div_or_zero(z.value, h), z.partials)
+    end,
+    begin
+        h = hypot(x.value, y, z)
+        dual_definition_retval(Val{Tx}(), h, _div_or_zero(x.value, h), x.partials)
+    end,
+    begin
+        h = hypot(x, y.value, z)
+        dual_definition_retval(Val{Ty}(), h, _div_or_zero(y.value, h), y.partials)
+    end,
+    begin
+        h = hypot(x, y, z.value)
+        dual_definition_retval(Val{Tz}(), h, _div_or_zero(z.value, h), z.partials)
+    end,
 )
 
+# For four or more arguments, Base promotes the arguments and calls `_hypot` with a tuple
+function Base.Math._hypot(x::DualNTupleRealComplex)
+    vx = _norm_primals(x)
+    return _norm_dual_x(x, vx, hypot(vx...), 2)
+end
+
+# LinearAlgebra.norm #
+#--------------------#
+
+# The generic rescaling branches incorrectly for Duals, so we differentiate the primal norm directly
+const DualArrayRealComplex{T,V,N} = Union{AbstractArray{Dual{T,V,N}},AbstractArray{Complex{Dual{T,V,N}}}}
+LinearAlgebra.norm(x::DualArrayRealComplex, p::Real=2) = _norm_dual_x(x, p)
+LinearAlgebra.norm(x::AbstractArray{<:Number}, p::Dual) = _norm_dual_p(x, p)
+LinearAlgebra.norm(x::DualArrayRealComplex, p::Dual) = _norm_dual_xp(x, p)
+
+# Tuples are promoted to a common element type
+const DualTuple = Tuple{Union{Dual,Complex{<:Dual}},Vararg{Union{Dual,Complex{<:Dual}}}}
+LinearAlgebra.norm(x::DualTuple, p::Real=2) = _norm_dual_x(promote(x...), p)
+LinearAlgebra.norm(x::Tuple{Number,Vararg{Number}}, p::Dual) = _norm_dual_p(x, p)
+LinearAlgebra.norm(x::DualTuple, p::Dual) = _norm_dual_xp(promote(x...), p)
+
+const DualContainerRealComplex{T,V,N} = Union{DualArrayRealComplex{T,V,N},DualNTupleRealComplex{T,V,N}}
+
+# Resolve method ambiguities
+const AdjOrTransDualRealComplex{T,V,N} = Union{
+    Adjoint{Dual{T,V,N}},Adjoint{Complex{Dual{T,V,N}}},Transpose{Dual{T,V,N}},Transpose{Complex{Dual{T,V,N}}}
+}
+for P in (:Real, :Dual)
+    @eval function LinearAlgebra.norm(x::AdjOrTransDualRealComplex, p::$P)
+        return invoke(norm, Tuple{Union{Adjoint,Transpose},Real}, x, p)
+    end
+end
+function LinearAlgebra.norm(x::Union{Adjoint{<:Number},Transpose{<:Number}}, p::Dual)
+    return invoke(norm, Tuple{Union{Adjoint,Transpose},Real}, x, p)
+end
+
+function _norm_dual_xp(x::DualContainerRealComplex{T}, p::Dual{T}) where {T}
+    vx = _norm_primals(x)
+    # Not `norm(x, p.value)`: it would confuse `p.value` with `p` if it is a Dual with tag `T` as well (#845)
+    y = _norm_dual_x(x, vx, norm(vx, p.value), p.value)
+    ∂p = _norm_deriv_p(vx, y.value, p.value)
+    return Dual{T}(y.value, y.partials + ∂p * p.partials)
+end
+# Arguments that are not Duals of the result tag are constants
+function _norm_dual_xp(x::DualContainerRealComplex{Tx}, p::Dual{Tp}) where {Tx,Tp}
+    return Tp ≺ Tx ? _norm_dual_x(x, p) : _norm_dual_p(x, p)
+end
+
+function _norm_dual_p(x::Union{AbstractArray{<:Number},Tuple{Number,Vararg{Number}}}, p::Dual{T}) where {T}
+    y = norm(x, p.value)
+    return Dual{T}(y, _norm_deriv_p(x, y, p.value) * p.partials)
+end
+
+# Uses `sum(u.^p) == 1` for `u = abs.(x) ./ norm(x, p)`
+function _norm_deriv_p(x, y, p)
+    if _iszero_primal(p) || isinf(p) || _iszero_primal(y)
+        return zero(y)
+    end
+    s = sum(x) do xi
+        u = abs(xi) / y
+        return _iszero_primal(u) ? zero(u) : u^p * log(u)
+    end
+    return oftype(y, y * s / p)
+end
+
+function _norm_dual_x(x::DualContainerRealComplex, p::Real)
+    # As in LinearAlgebra; also avoids relying on `map` inferring the element type for empty arrays
+    if isempty(x)
+        return float(norm(zero(eltype(x))))
+    end
+    vx = _norm_primals(x)
+    return _norm_dual_x(x, vx, norm(vx, p), p)
+end
+# `Partials` are `AbstractVector`s, but we work with their tuple of values
+_norm_dual_x(x::Partials{N}, p::Real) where {N} = N == 0 ? float(norm(zero(eltype(x)))) : _norm_dual_x(x.values, p)
+
+# Partials of `y = norm(x, p)`, given the primal values `vx` of `x`
+function _norm_dual_x(x::DualContainerRealComplex{T,V,N}, vx, y, p::Real) where {T,V,N}
+    if _iszero_primal(p) || _iszero_primal(y)
+        return Dual{T}(y, zero(Partials{N,typeof(y)}))
+    end
+    ∂y = if p == 2
+        _sum_map((xi, vxi) -> _norm_partials(vxi / y, xi), x, vx)
+    elseif isinf(p)
+        # First element with absolute value `y` (faster than `findmax`/`findmin`)
+        i = p > 0 ? findfirst(v -> !(abs(v) < y), vx) : findfirst(v -> !(abs(v) > y), vx)
+        _norm_partials(vx[i] / y, x[i])
+    else
+        _sum_map(x, vx) do xi, vxi
+            r = vxi / y
+            w = _iszero_primal(r) ? zero(r) : oftype(r, abs(r)^(p - 1) * sign(r))
+            return _norm_partials(w, xi)
+        end
+    end
+    return Dual{T}(y, ∂y)
+end
+
+_norm_primal(x::Dual) = x.value
+_norm_primal(z::Complex{<:Dual}) = complex(real(z).value, imag(z).value)
+
+# Primal values; dense arrays of real Duals are reinterpreted without copying (strided, e.g., for BLAS)
+_norm_primals(x) = map(_norm_primal, x)
+function _norm_primals(x::DenseArray{Dual{T,V,N}}) where {T,V,N}
+    if !isbitstype(V)
+        return map(_norm_primal, x)
+    elseif N == 0
+        # `reinterpret(reshape, V, x)` has the same size as `x`
+        return reinterpret(V, x)
+    else
+        return selectdim(reinterpret(reshape, V, x), 1, 1)
+    end
+end
+
+# Partials of `real(conj(w) * x)`
+_norm_partials(w::Real, x::Dual) = w * x.partials
+_norm_partials(w::Complex, x::Complex{<:Dual}) = _mul_partials(real(x).partials, imag(x).partials, real(w), imag(w))
+
+# `sum(map(f, x, vx))` without allocating
+_sum_map(f::F, x::Tuple, vx::Tuple) where {F} = sum(map(f, x, vx))
+_sum_map(f::F, x::AbstractArray, vx::AbstractArray) where {F} = sum(Broadcast.instantiate(Broadcast.broadcasted(f, x, vx)))
+
 # fma #
 #-----#
 

test/DualTest.jl

@@ -775,4 +775,51 @@ end
     end
 end
 
+@testset "hypot" begin
+    @testset "zero" begin
+        x = Dual{TestTag}(0.0, 1.0, -2.0)
+        y = Dual{TestTag}(-0.0, 3.0, 0.5)
+        for h in (
+            hypot(x, y), hypot(x, 0.0), hypot(0.0, y), hypot(x, y, x), hypot(x, 0.0, y), abs(complex(x, y)),
+            hypot(x, y, x, y), hypot(0.0, x, 0, y), hypot(complex(x, y), x, y), hypot(complex(x, y), 0.0, x, 0),
+        )
+            @test h == Dual{TestTag}(0.0, 0.0, 0.0)
+        end
+    end
+
+    @testset "infinite" begin
+        x = Dual{TestTag}(Inf, 0.0)
+        y = Dual{TestTag}(1.0, 1.0)
+        for h in (hypot(x, y), hypot(x, y, y), hypot(x, y, y, y))
+            if ForwardDiff.NANSAFE_MODE_ENABLED
+                @test h === Dual{TestTag}(Inf, 0.0)
+            else
+                @test isequal(h, Dual{TestTag}(Inf, NaN))
+            end
+        end
+    end
+
+    @testset "nested differentiation" begin
+        a, b = 3.0, 4.0
+        # Mixed tags: the outer variable `a` and the inner variable `b`
+        for (f, g) in (
+            ((a, b) -> hypot(a, b), (a, b) -> sqrt(a^2 + b^2)),
+            ((a, b) -> hypot(a, a, b), (a, b) -> sqrt(2a^2 + b^2)),
+            ((a, b) -> hypot(a, b, a), (a, b) -> sqrt(2a^2 + b^2)),
+            ((a, b) -> hypot(b, a, a), (a, b) -> sqrt(2a^2 + b^2)),
+            ((a, b) -> hypot(a, b, b), (a, b) -> sqrt(a^2 + 2b^2)),
+            ((a, b) -> hypot(b, a, b), (a, b) -> sqrt(a^2 + 2b^2)),
+            ((a, b) -> hypot(b, b, a), (a, b) -> sqrt(a^2 + 2b^2)),
+            ((a, b) -> hypot(a, 2.0, b), (a, b) -> sqrt(a^2 + 4 + b^2)),
+            ((a, b) -> hypot(a, b, a, 2.0), (a, b) -> sqrt(2a^2 + b^2 + 4)),
+            ((a, b) -> hypot(b, 2.0, a, b, a), (a, b) -> sqrt(2a^2 + 2b^2 + 4)),
+            ((a, b) -> hypot(complex(a, b), a, b), (a, b) -> sqrt(2a^2 + 2b^2)),
+        )
+            df = ForwardDiff.derivative(a -> ForwardDiff.derivative(b -> f(a, b), b), a)
+            dg = ForwardDiff.derivative(a -> ForwardDiff.derivative(b -> g(a, b), b), a)
+            @test df ≈ dg
+        end
+    end
+end
+
 end # module