From 50ba6b3da508bde82ffe4a4e82bb118dea4733e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 4 Nov 2026 01:02:42 +0100 Subject: [PATCH 1/2] Add derivative rules for `norm` and fix `hypot` at zero The generic implementations of `norm` in LinearAlgebra and StaticArrays rescale by the maximum absolute value to avoid overflow and underflow. Their branches (`iszero`, `isfinite`, `<`) take partials into account for `Dual`s, resulting in NaNs, Infs, and even incorrect primal values (e.g., at zero or for very small or large inputs). Instead, the norm of the primal values is computed and the chain rule is applied with weights that are bounded by 1 in absolute value. At zero, the subgradient zero is used (as in ChainRules). Derivatives with respect to `p` are supported as well. Additionally, `hypot` uses the same convention at zero, and the three-argument `hypot` no longer drops the partials of outer tags in nested differentiation. Fixes #243, fixes #794. Co-Authored-By: Claude Opus 5.5 --- Project.toml | 3 + docs/src/user/limitations.md | 2 +- ext/ForwardDiffStaticArraysExt.jl | 5 + ext/SparseArraysExt.jl | 22 +++++ src/dual.jl | 146 ++++++++++++++++++++++++++---- test/DualTest.jl | 29 ++++++ test/GradientTest.jl | 64 +++++++++++++ 7 files changed, 254 insertions(+), 17 deletions(-) create mode 100644 ext/SparseArraysExt.jl diff --git a/Project.toml b/Project.toml index 396b3f7f..5cf74ee9 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/user/limitations.md b/docs/src/user/limitations.md index f90de787..f9ceac56 100644 --- a/docs/src/user/limitations.md +++ b/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. diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index 796bdf5c..7e6c69fa 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -42,6 +42,11 @@ 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) + # For `MMatrix` we can use the in-place method ForwardDiff._lyap_div_zero_diag!!(A::StaticArrays.MMatrix, λ::AbstractVector) = ForwardDiff._lyap_div_zero_diag!(A, λ) diff --git a/ext/SparseArraysExt.jl b/ext/SparseArraysExt.jl new file mode 100644 index 00000000..d0e260de --- /dev/null +++ b/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 diff --git a/src/dual.jl b/src/dual.jl index 48d5d0e3..38f61d58 100644 --- a/src/dual.jl +++ b/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,140 @@ 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 + +# 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, ) +# 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) + +# 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::DualArrayRealComplex{T}, p::Dual{T}) where {T} + y = norm(x, p.value) + ∂p = _norm_deriv_p(map(_norm_primal, x), 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::DualArrayRealComplex{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::AbstractArray{<: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) + (_iszero_primal(p) || isinf(p) || _iszero_primal(y)) && return zero(y / p) + s = sum(x) do xi + u = abs(xi) / y + return _iszero_primal(u) ? zero(u) : u^p * log(u) + end + return y * s / p +end + +function _norm_dual_x(x::DualArrayRealComplex{T,V,N}, p::Real) where {T,V,N} + # These cases are not affected by the scaling issues + (p == 0 || p == 1 || p == -Inf) && return invoke(norm, Tuple{Any,Real}, x, p) + isempty(x) && return float(norm(zero(eltype(x)))) + vx = map(_norm_primal, x) + y = norm(vx, p) + ∂y = zero(Partials{N,typeof(y)}) + _iszero_primal(y) && return Dual{T}(y, ∂y) + if p == 2 + for (xi, vxi) in zip(x, vx) + ∂y += _norm_partials(vxi / y, xi) + end + elseif p == Inf + _, i = findmax(abs, vx) + ∂y = _norm_partials(vx[i] / y, x[i]) + else + for (xi, vxi) in zip(x, vx) + r = vxi / y + w = _iszero_primal(r) ? zero(r) : abs(r)^(p - 1) * sign(r) + ∂y += _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) + +# 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)) + # fma # #-----# diff --git a/test/DualTest.jl b/test/DualTest.jl index e7429be5..44fa6809 100644 --- a/test/DualTest.jl +++ b/test/DualTest.jl @@ -775,4 +775,33 @@ 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))) + @test h == Dual{TestTag}(0.0, 0.0, 0.0) + 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)), + ) + 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 diff --git a/test/GradientTest.jl b/test/GradientTest.jl index bf121239..f6d9fe83 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -9,6 +9,7 @@ using LinearAlgebra using ForwardDiff using ForwardDiff: Dual, Tag using StaticArrays +using SparseArrays: sparse using DiffTests include(joinpath(dirname(@__FILE__), "utils.jl")) @@ -347,4 +348,67 @@ end end end +@testset "norm" begin + # issues #243 and #794 + for x in ([0.0, 0.0], SVector(0.0, 0.0)), p in (2, 3, 1.5, Inf) + @test ForwardDiff.gradient(x -> norm(x, p), x) == [0.0, 0.0] + @test ForwardDiff.gradient(x -> norm(x, p)^2, x) == [0.0, 0.0] + end + @test norm([Dual(0.0, 1.0), Dual(0.0, 1.0)]) === Dual(0.0, 0.0) + @test norm(SVector(Dual(0.0, 1.0), Dual(0.0, 1.0), Dual(0.0, 1.0))) === Dual(0.0, 0.0) + @test norm([Dual(1e-200, 1e-100), Dual(0.0, 0.0)]) === Dual(1e-200, 1e-100) + for s in (1e-200, 1e200) + @test ForwardDiff.gradient(norm, [s, s]) ≈ [√0.5, √0.5] + @test ForwardDiff.gradient(norm, SVector(s, s)) ≈ [√0.5, √0.5] + end + @test ForwardDiff.gradient(x -> norm(x, 1.01), [5e-324, 1.0]) ≈ [5e-324^0.01, 1.0] + + # `p` is a `Dual` + ref = (x, p) -> sum(t -> abs(t)^p, x)^(1 / p) + x = [3.0, -4.0, 0.5] + for p in (3.0, 1.5, 0.5) + dref = ForwardDiff.derivative(p -> ref(x, p), p) + @test ForwardDiff.derivative(p -> norm(x, p), p) ≈ dref + @test ForwardDiff.derivative(p -> norm(SVector{3}(x), p), p) ≈ dref + @test ForwardDiff.derivative(p -> norm(x', p), p) ≈ dref + @test ForwardDiff.derivative(p -> norm(sparse(x), p), p) ≈ dref + @test ForwardDiff.derivative(p -> norm(complex.(x, x), p), p) ≈ ForwardDiff.derivative(p -> ref(complex.(x, x), p), p) + @test ForwardDiff.gradient(v -> norm(v[1:3], v[4]), [x; p]) ≈ ForwardDiff.gradient(v -> ref(v[1:3], v[4]), [x; p]) + @test ForwardDiff.derivative(p -> ForwardDiff.gradient(x -> norm(x, p), x), p) ≈ + ForwardDiff.derivative(p -> ForwardDiff.gradient(x -> ref(x, p), x), p) + @test ForwardDiff.gradient(x -> ForwardDiff.derivative(p -> norm(x, p), p), x) ≈ + ForwardDiff.gradient(x -> ForwardDiff.derivative(p -> ref(x, p), p), x) + end + @test ForwardDiff.derivative(p -> norm([0.0, 0.0], p), 3.0) == 0 + @test ForwardDiff.derivative(p -> norm([1e200, 1e200], p), 3.0) ≈ 1e200 * ForwardDiff.derivative(p -> ref([1.0, 1.0], p), 3.0) + + x = [1.0, -2.0, 0.0, 3.0] + for p in (2, 3, 1.5, 1, Inf) + f = x -> norm(x, p) + g = p == Inf ? (x -> maximum(abs, x)) : (x -> sum(t -> abs(t)^p, x)^(1 / p)) + ∇g = ForwardDiff.gradient(g, x) + for c in (1, 4) + @test ForwardDiff.gradient(f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}())) ≈ ∇g + end + @test ForwardDiff.gradient(f, SVector{4}(x)) ≈ ∇g + + z = v -> complex.(v[1:2], v[3:4]) + @test ForwardDiff.gradient(f ∘ z, x) ≈ ForwardDiff.gradient(g ∘ z, x) + end + + u = x / norm(x) + @test ForwardDiff.hessian(norm, x) ≈ (I - u * u') / norm(x) + @test ForwardDiff.hessian(norm, SVector{4}(x)) ≈ (I - u * u') / norm(x) + @test ForwardDiff.hessian(x -> norm(x, 3), x) ≈ ForwardDiff.hessian(x -> cbrt(sum(t -> abs(t)^3, x)), x) + + d = [Dual(3.0, 1.0), Dual(-4.0, 2.0)] + for p in (2, 3, 1, Inf) + y = norm(d, p) + @test norm(d', p) === y + @test norm(transpose(d), p) === y + @test norm(sparse(d), p) === y + @test norm(sparse(reshape(d, 2, 1)), p) === y + end +end + end # module From a390d78d15280cc442b85a7ad0106df6ffd17b63 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 4 Nov 2026 16:02:45 +0100 Subject: [PATCH 2/2] Support tuples and n-ary `hypot`, fix mixed Hessians and p = 0 in `norm` - Compute the partials of `norm(x, p)` with respect to `x` for all `p` in one code path: the derivative is zero for `p = 0` (whose primal value was incorrect), only the first element with the largest (smallest) absolute value contributes for `p = Inf` (`p = -Inf`), and zero elements get the subgradient zero (also for `p = 1`). - Compute the derivative factors in the precision of the primal norm, which fixes a type instability for Float32 inputs. - Do not dispatch on `p.value` when differentiating with respect to both `x` and `p`: with the same tag at multiple levels (as in `hessian`), it was confused with `p` (#845). - Support tuples (promoted to a common element type) and `hypot` with four or more arguments, and use the tuple of values of `Partials`. - Reinterpret dense arrays of Duals without copying (strided, e.g., for BLAS), and sum the partials without allocating. Co-Authored-By: Claude Opus 5.5 --- ext/ForwardDiffStaticArraysExt.jl | 2 + src/dual.jl | 96 +++++++++++++++++++++++-------- test/DualTest.jl | 20 ++++++- test/GradientTest.jl | 56 ++++++++++++++++-- 4 files changed, 143 insertions(+), 31 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index 7e6c69fa..807bc589 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -46,6 +46,8 @@ const StaticDualArrayRealComplex{T,V,N} = Union{StaticArray{<:Tuple,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, λ) diff --git a/src/dual.jl b/src/dual.jl index 38f61d58..7963b9bc 100644 --- a/src/dual.jl +++ b/src/dual.jl @@ -604,6 +604,10 @@ _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, @@ -654,6 +658,12 @@ _div_or_zero(v, h) = _iszero_primal(h) ? zero(v / oneunit(h)) : v / h 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 # #--------------------# @@ -663,6 +673,14 @@ 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}}} @@ -676,51 +694,62 @@ function LinearAlgebra.norm(x::Union{Adjoint{<:Number},Transpose{<:Number}}, p:: return invoke(norm, Tuple{Union{Adjoint,Transpose},Real}, x, p) end -function _norm_dual_xp(x::DualArrayRealComplex{T}, p::Dual{T}) where {T} - y = norm(x, p.value) - ∂p = _norm_deriv_p(map(_norm_primal, x), y.value, p.value) +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::DualArrayRealComplex{Tx}, p::Dual{Tp}) where {Tx,Tp} +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::AbstractArray{<:Number}, p::Dual{T}) where {T} +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) - (_iszero_primal(p) || isinf(p) || _iszero_primal(y)) && return zero(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 y * s / p + return oftype(y, y * s / p) end -function _norm_dual_x(x::DualArrayRealComplex{T,V,N}, p::Real) where {T,V,N} - # These cases are not affected by the scaling issues - (p == 0 || p == 1 || p == -Inf) && return invoke(norm, Tuple{Any,Real}, x, p) - isempty(x) && return float(norm(zero(eltype(x)))) - vx = map(_norm_primal, x) - y = norm(vx, p) - ∂y = zero(Partials{N,typeof(y)}) - _iszero_primal(y) && return Dual{T}(y, ∂y) - if p == 2 - for (xi, vxi) in zip(x, vx) - ∂y += _norm_partials(vxi / y, xi) - end - elseif p == Inf - _, i = findmax(abs, vx) - ∂y = _norm_partials(vx[i] / y, x[i]) +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 - for (xi, vxi) in zip(x, vx) + _sum_map(x, vx) do xi, vxi r = vxi / y - w = _iszero_primal(r) ? zero(r) : abs(r)^(p - 1) * sign(r) - ∂y += _norm_partials(w, xi) + 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) @@ -729,10 +758,27 @@ 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 # #-----# diff --git a/test/DualTest.jl b/test/DualTest.jl index 44fa6809..5cf12afe 100644 --- a/test/DualTest.jl +++ b/test/DualTest.jl @@ -779,11 +779,26 @@ end @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))) + 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` @@ -796,6 +811,9 @@ end ((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) diff --git a/test/GradientTest.jl b/test/GradientTest.jl index f6d9fe83..b0b814c9 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -350,13 +350,38 @@ end @testset "norm" begin # issues #243 and #794 - for x in ([0.0, 0.0], SVector(0.0, 0.0)), p in (2, 3, 1.5, Inf) + for x in ([0.0, 0.0], SVector(0.0, 0.0)), p in (2, 3, 1.5, 1, 0, Inf, -Inf) @test ForwardDiff.gradient(x -> norm(x, p), x) == [0.0, 0.0] @test ForwardDiff.gradient(x -> norm(x, p)^2, x) == [0.0, 0.0] + @test ForwardDiff.gradient(x -> norm(Tuple(x), p), x) == [0.0, 0.0] end @test norm([Dual(0.0, 1.0), Dual(0.0, 1.0)]) === Dual(0.0, 0.0) @test norm(SVector(Dual(0.0, 1.0), Dual(0.0, 1.0), Dual(0.0, 1.0))) === Dual(0.0, 0.0) + @test norm((Dual(0.0, 1.0), Dual(0.0, 1.0))) === Dual(0.0, 0.0) + @test norm((Dual(0.0, 1.0), complex(Dual(0.0, 1.0), Dual(0.0, 1.0))), 3) === Dual(0.0, 0.0) + @test norm(ForwardDiff.partials(Dual(Dual(1.0, 0.0), Dual(0.0, 1.0), Dual(0.0, 1.0)))) === Dual(0.0, 0.0) @test norm([Dual(1e-200, 1e-100), Dual(0.0, 0.0)]) === Dual(1e-200, 1e-100) + for x in ([Dual(0.0, 1.0), Dual(1.0, 0.0)], SVector(Dual(0.0, 1.0), Dual(1.0, 0.0)), (Dual(0.0, 1.0), Dual(1.0, 0.0))) + @test norm(x, 0) === Dual(1.0, 0.0) + end + + # Infinite entries yield `NaN` partials, unless NaN-safe mode is enabled + for x in ([Dual(Inf, 0.0), Dual(1.0, 1.0)], (Dual(Inf, 0.0), Dual(1.0, 1.0))), p in (2, 3, Inf) + if ForwardDiff.NANSAFE_MODE_ENABLED + @test norm(x, p) === Dual(Inf, 0.0) + else + @test isequal(norm(x, p), Dual(Inf, NaN)) + end + end + + # Same precision as the primal norm + for x in ([Dual(1f0, 1f0), Dual(-2f0, 0f0)], SVector(Dual(1f0, 1f0), Dual(-2f0, 0f0)), (Dual(1f0, 1f0), Dual(-2f0, 0f0))), p in (2, 3, 2.5, Inf) + y = norm(map(ForwardDiff.value, x), p) + @test norm(x, p) isa Dual{Nothing,typeof(y),1} + if !(x isa StaticArray) + @test @inferred(norm(x, p)) isa Dual{Nothing,Float32,1} + end + end for s in (1e-200, 1e200) @test ForwardDiff.gradient(norm, [s, s]) ≈ [√0.5, √0.5] @test ForwardDiff.gradient(norm, SVector(s, s)) ≈ [√0.5, √0.5] @@ -372,18 +397,34 @@ end @test ForwardDiff.derivative(p -> norm(SVector{3}(x), p), p) ≈ dref @test ForwardDiff.derivative(p -> norm(x', p), p) ≈ dref @test ForwardDiff.derivative(p -> norm(sparse(x), p), p) ≈ dref + @test ForwardDiff.derivative(p -> norm(sparse(reshape(x, 3, 1)), p), p) ≈ dref + @test ForwardDiff.derivative(p -> norm(Tuple(x), p), p) ≈ dref @test ForwardDiff.derivative(p -> norm(complex.(x, x), p), p) ≈ ForwardDiff.derivative(p -> ref(complex.(x, x), p), p) - @test ForwardDiff.gradient(v -> norm(v[1:3], v[4]), [x; p]) ≈ ForwardDiff.gradient(v -> ref(v[1:3], v[4]), [x; p]) + ∇ref = ForwardDiff.gradient(v -> ref(v[1:3], v[4]), [x; p]) + @test ForwardDiff.gradient(v -> norm(v[1:3], v[4]), [x; p]) ≈ ∇ref + @test ForwardDiff.gradient(v -> norm(SVector{3}(v[1], v[2], v[3]), v[4]), [x; p]) ≈ ∇ref + @test ForwardDiff.gradient(v -> norm(v[1:3]', v[4]), [x; p]) ≈ ∇ref + @test ForwardDiff.gradient(v -> norm(sparse(v[1:3]), v[4]), [x; p]) ≈ ∇ref + @test ForwardDiff.gradient(v -> norm(sparse(reshape(v[1:3], 3, 1)), v[4]), [x; p]) ≈ ∇ref + @test ForwardDiff.gradient(v -> norm((v[1], v[2], v[3]), v[4]), [x; p]) ≈ ∇ref @test ForwardDiff.derivative(p -> ForwardDiff.gradient(x -> norm(x, p), x), p) ≈ ForwardDiff.derivative(p -> ForwardDiff.gradient(x -> ref(x, p), x), p) @test ForwardDiff.gradient(x -> ForwardDiff.derivative(p -> norm(x, p), p), x) ≈ ForwardDiff.gradient(x -> ForwardDiff.derivative(p -> ref(x, p), p), x) end + for p in (1.0, 2.0, 3.0) + Href = ForwardDiff.hessian(v -> ref(v[1:3], v[4]), [x; p]) + @test ForwardDiff.hessian(v -> norm(v[1:3], v[4]), [x; p]) ≈ Href + @test ForwardDiff.hessian(v -> norm(v[SOneTo(3)], v[4]), SVector{4}([x; p])) ≈ Href + @test ForwardDiff.hessian(v -> norm((v[1], v[2], v[3]), v[4]), [x; p]) ≈ Href + end + @test ForwardDiff.derivative(p -> ForwardDiff.gradient(x -> norm(x, p), x), Inf) == zeros(3) @test ForwardDiff.derivative(p -> norm([0.0, 0.0], p), 3.0) == 0 @test ForwardDiff.derivative(p -> norm([1e200, 1e200], p), 3.0) ≈ 1e200 * ForwardDiff.derivative(p -> ref([1.0, 1.0], p), 3.0) x = [1.0, -2.0, 0.0, 3.0] - for p in (2, 3, 1.5, 1, Inf) + z = v -> complex.(v[1:2], v[3:4]) + for p in (2, 3, 1.5, Inf) f = x -> norm(x, p) g = p == Inf ? (x -> maximum(abs, x)) : (x -> sum(t -> abs(t)^p, x)^(1 / p)) ∇g = ForwardDiff.gradient(g, x) @@ -391,10 +432,14 @@ end @test ForwardDiff.gradient(f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}())) ≈ ∇g end @test ForwardDiff.gradient(f, SVector{4}(x)) ≈ ∇g - - z = v -> complex.(v[1:2], v[3:4]) + @test ForwardDiff.gradient(f ∘ Tuple, x) ≈ ∇g @test ForwardDiff.gradient(f ∘ z, x) ≈ ForwardDiff.gradient(g ∘ z, x) end + # Subgradient zero for zero entries, as for other `p` + @test ForwardDiff.gradient(x -> norm(x, 1), x) == sign.(x) + @test ForwardDiff.gradient(x -> norm(x, 1), SVector{4}(x)) == sign.(x) + @test ForwardDiff.gradient(x -> norm(z(x), 1), x) ≈ ForwardDiff.gradient(x -> sum(abs, z(x)), x) + @test ForwardDiff.gradient(x -> norm(x, -Inf), [3.0, -1.0, 2.0]) == [0.0, -1.0, 0.0] u = x / norm(x) @test ForwardDiff.hessian(norm, x) ≈ (I - u * u') / norm(x) @@ -408,6 +453,7 @@ end @test norm(transpose(d), p) === y @test norm(sparse(d), p) === y @test norm(sparse(reshape(d, 2, 1)), p) === y + @test norm(Tuple(d), p) === y end end