Enable NaN-safe mode by default

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ permissions:
 
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - NaN-safe mode ${{ matrix.nansafe && 'enabled' || 'disabled' }} - ${{ github.event_name }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - unsafe mode ${{ matrix.unsafe && 'enabled' || 'disabled' }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -29,7 +29,7 @@ jobs:
           - ubuntu-latest
           - windows-latest
           - macOS-latest
-        nansafe:
+        unsafe:
           - true
           - false
     steps:
@@ -38,10 +38,10 @@ jobs:
         with:
           version: ${{ matrix.version }}
       - uses: julia-actions/cache@v3
-      - name: Set NaN-safe mode preference
+      - name: Set unsafe mode preference
         run: |
           open(joinpath(pwd(), "Project.toml"), "a") do io
-            println(io, "\n[preferences.ForwardDiff]\nnansafe_mode = ${{ matrix.nansafe }}\n")
+            println(io, "\n[preferences.ForwardDiff]\nunsafe_mode = ${{ matrix.unsafe }}\n")
           end
         shell: julia --color=yes {0}
       - uses: julia-actions/julia-buildpkg@v1

Project.toml

@@ -1,6 +1,6 @@
 name = "ForwardDiff"
 uuid = "f6369f11-7733-5829-9624-2563aa707210"
-version = "1.4.6"
+version = "1.5.0"
 
 [deps]
 CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"

docs/src/user/advanced.md

@@ -117,41 +117,36 @@ Note that it is usually best to pick a chunk size which divides evenly into the
 dimension. Otherwise, ForwardDiff has to construct and utilize an extra "remainder" chunk to
 complete the calculation.
 
-## Fixing NaN/Inf Issues
+## NaN-Safe Mode
 
-ForwardDiff's default behavior is to return `NaN` for undefined derivatives (or otherwise
-mirror the behavior of the function in `Base`, if it would return an error). This is
-usually the correct thing to do, but in some cases can erroneously "poison" values which
-aren't sensitive to the input and thus cause ForwardDiff to incorrectly return `NaN` or
-`Inf` derivatives. For example:
+A dual number whose perturbation component is zero is unperturbed: it must not propagate
+derivative information, regardless of the derivative of the function it is passed to.
+Naively propagating derivatives by multiplying them with the perturbation component
+violates this property whenever the derivative is not finite, since `0.0 * Inf` and
+`0.0 * NaN` are `NaN`. ForwardDiff therefore checks whether each perturbation component is
+zero before propagating derivative information ("`NaN`-safe mode"):
 
 ```julia-repl
-# the dual number's perturbation component is zero, so this
-# variable should not propagate derivative information
 julia> log(ForwardDiff.Dual{:tag}(0.0, 0.0))
-Dual{:tag}(-Inf,NaN) # oops, this NaN should be 0.0
-```
-
-Here, ForwardDiff computes the derivative of `log(0.0)` as `NaN` and then propagates
-this derivative by multiplying it by the perturbation component. Usually, ForwardDiff can
-rely on the identity `x * 0.0 == 0.0` to prevent the derivatives from propagating when
-the perturbation component is `0.0`. However, this identity doesn't hold if `isnan(y)`
-or `isinf(y)`, in which case a `NaN` derivative will be propagated instead.
+Dual{:tag}(-Inf,0.0)
 
-It is possible to fix this behavior by checking that the perturbation component is zero
-before attempting to propagate derivative information, but this check can noticeably
-decrease performance (~5%-10% on our benchmarks).
+julia> ForwardDiff.gradient(v -> sqrt(v[1]) + v[2], [0.0, 1.0])
+2-element Vector{Float64}:
+ Inf
+  1.0
+```
 
-In order to preserve performance in the majority of use cases, ForwardDiff disables this
-check by default. If your code is affected by this `NaN` behavior, you can enable
-ForwardDiff's `NaN`-safe mode by using the
+Similar to `@fastmath`, you can trade correctness for performance by disabling `NaN`-safe
+mode: derivatives may then be `NaN` even though they are well-defined (see the example
+below). Unsafe mode can make code faster, in particular code that is dominated by cheap
+arithmetic operations. To enable it, use the
 [Preferences.jl](https://github.com/JuliaPackaging/Preferences.jl) API to set
-the `nansafe_mode` preference to true, for example via:
+the `unsafe_mode` preference to true, for example via:
 
 ```julia-repl
 julia> using ForwardDiff, Preferences
 
-julia> set_preferences!(ForwardDiff, "nansafe_mode" => true)
+julia> set_preferences!(ForwardDiff, "unsafe_mode" => true)
 ```
 
 Note that Julia has to be restarted and ForwardDiff has to be reloaded after changing
@@ -162,16 +157,27 @@ Alternatively, you can set the preference before loading ForwardDiff, for exampl
 ```julia-repl
 julia> using Preferences, UUIDs
 
-julia> set_preferences!(UUID("f6369f11-7733-5829-9624-2563aa707210"), "nansafe_mode" => true)
+julia> set_preferences!(UUID("f6369f11-7733-5829-9624-2563aa707210"), "unsafe_mode" => true)
 
 julia> using ForwardDiff
 
 julia> log(ForwardDiff.Dual{:tag}(0.0, 0.0))
-Dual{:tag}(-Inf,0.0)
+Dual{:tag}(-Inf,NaN)
+
+julia> ForwardDiff.gradient(v -> sqrt(v[1]) + v[2], [0.0, 1.0])
+2-element Vector{Float64}:
+  Inf
+ NaN
 ```
 
+!!! compat "ForwardDiff 1.5"
+    Before ForwardDiff 1.5.0, `NaN`-safe mode was disabled by default and could be enabled by
+    setting the `nansafe_mode` preference to true. The `nansafe_mode` preference is
+    deprecated: setting it emits a warning, and `nansafe_mode = false` disables `NaN`-safe
+    mode only if the `unsafe_mode` preference is not set.
+
 In the future, we plan on allowing users and downstream library authors to dynamically
-enable [`NaN`-safe mode via the `AbstractConfig`
+toggle [`NaN`-safe mode via the `AbstractConfig`
 API](https://github.com/JuliaDiff/ForwardDiff.jl/issues/181).
 
 ## Hessian of a vector-valued function

src/ForwardDiff.jl

@@ -24,4 +24,18 @@ include("hessian.jl")
 
 export DiffResults
 
+function __init__()
+    if NANSAFE_MODE_PREFERENCE !== nothing
+        Base.depwarn(
+            "The `nansafe_mode` preference of ForwardDiff is deprecated. " *
+            "NaN-safe mode is enabled by default since ForwardDiff 1.5.0, and the preference to disable it is now called `unsafe_mode`. " *
+            "For now, `nansafe_mode = false` still disables NaN-safe mode unless `unsafe_mode` is set. " *
+            "Please remove the `nansafe_mode` preference with `Preferences.delete_preferences!(ForwardDiff, \"nansafe_mode\"; force = true)` and restart Julia.",
+            :__init__;
+            force = true,
+        )
+    end
+    return nothing
+end
+
 end # module

src/dual.jl

@@ -416,6 +416,10 @@ end
 
 Base.iszero(x::Dual) = iszero(value(x)) && iszero(partials(x))  # shortcut, equivalent to x == zero(x)
 
+@inline _isfinite_all(x::Dual) = _isfinite_all(x.value) & all(_isfinite_all, x.partials.values)
+@inline _is_strong_divisor(x::Dual) = _is_strong_divisor(x.value) & all(_isfinite_all, x.partials.values)
+@inline _is_strong_zero_type(::Type{<:Dual{<:Any,V}}) where {V} = _is_strong_zero_type(V)
+
 for pred in [:isequal, :(==)]
     @eval begin
         @define_binary_dual_op(

src/partials.jl

@@ -84,44 +84,71 @@ Base.convert(::Type{Partials{N,V}}, partials::Partials{N,V}) where {N,V} = parti
 @inline Base.:-(partials::Partials) = Partials(minus_tuple(partials.values))
 @inline Base.:*(x::Real, partials::Partials) = partials*x
 
-@inline function Base.:*(partials::Partials, x::Real)
-    return Partials(scale_tuple(partials.values, x))
+@inline function Base.:*(partials::Partials{N,V}, x::Real) where {N,V}
+    if _mul_is_strong(V, x)
+        return Partials(scale_tuple(*, partials.values, x))
+    else
+        return Partials(scale_tuple(_strong_mul, partials.values, x))
+    end
 end
 
-@inline function Base.:/(partials::Partials, x::Real)
-    return Partials(div_tuple_by_scalar(partials.values, x))
+@inline function Base.:/(partials::Partials{N,V}, x::Real) where {N,V}
+    if _div_is_strong(V, x)
+        return Partials(scale_tuple(/, partials.values, x))
+    else
+        return Partials(scale_tuple(_strong_div, partials.values, x))
+    end
 end
 
-@inline function _mul_partials(a::Partials{N}, b::Partials{N}, x_a, x_b) where N
-    return Partials(mul_tuples(a.values, b.values, x_a, x_b))
+@inline function _mul_partials(a::Partials{N,A}, b::Partials{N,B}, x_a, x_b) where {N,A,B}
+    if _mul_is_strong(A, x_a) & _mul_is_strong(B, x_b)
+        return Partials(mul_tuples(*, a.values, b.values, x_a, x_b))
+    else
+        return Partials(mul_tuples(_strong_mul, a.values, b.values, x_a, x_b))
+    end
 end
 
+# Divide by `bval` instead of multiplying by `inv(bval)` or dividing by `bval^2`, which may overflow or underflow
 @inline function _div_partials(a::Partials, b::Partials, aval, bval)
-    return _mul_partials(a, b, inv(bval), -(aval / (bval*bval)))
+    return (a - (aval / bval) * b) / bval
 end
 
 # NaN/Inf-safe methods #
 #----------------------#
 
-if NANSAFE_MODE_ENABLED
-    # A dual number with a zero partial is just an unperturbed non-dual number
-    # Hence when propagated the resulting dual number is unperturbed as well,
-    # ie., its partial is zero as well, regardless of the primal value
-    # However, standard floating point multiplication/division would return `NaN`
-    # if the primal is not-finite/zero
-    @inline function _mul_partial(partial::Real, x::Real)
-        y = partial * x
-        return iszero(partial) ? zero(y) : y
+# Zero perturbations in input yield zero perturbations in output
+@inline function _strong_mul(partial::Real, x::Real)
+    y = partial * x
+    return iszero(partial) ? zero(y) : y
+end
+@inline function _strong_div(partial::Real, x::Real)
+    y = partial / x
+    return iszero(partial) ? zero(y) : y
+end
+
+# Whether plain `*`/`/` already treat zero partials of type `V` as strong zeros
+if !UNSAFE_MODE_ENABLED
+    @inline function _mul_is_strong(::Type{V}, x::Real) where {V}
+        return _is_strong_zero_type(V) && _is_strong_zero_type(typeof(x)) && _isfinite_all(x)
     end
-    @inline function _div_partial(partial::Real, x::Real)
-        y = partial / x
-        return iszero(partial) ? zero(y) : y
+    @inline function _div_is_strong(::Type{V}, x::Real) where {V}
+        return _is_strong_zero_type(V) && _is_strong_zero_type(typeof(x)) && _is_strong_divisor(x)
     end
 else
-    @inline _mul_partial(partial::Real, x::Real) = partial * x
-    @inline _div_partial(partial::Real, x::Real) = partial / x
+    @inline _mul_is_strong(::Type, ::Real) = true
+    @inline _div_is_strong(::Type, ::Real) = true
 end
 
+# Types whose zeros stay zero when multiplied by finite values of such types
+@inline _is_strong_zero_type(::Type) = false
+@inline _is_strong_zero_type(::Type{<:Union{Base.IEEEFloat,Base.BitInteger,Bool,BigInt,BigFloat,Rational{<:Union{Base.BitInteger,Bool,BigInt}},AbstractIrrational}}) = true
+
+# Whether `x` and, for dual numbers, all of its partials are finite
+@inline _isfinite_all(x::Real) = isfinite(x)
+
+# Whether zero divided by `x` is zero
+@inline _is_strong_divisor(x::Real) = !(isnan(x) | iszero(x))
+
 # edge cases where N == 0 #
 #-------------------------#
 
@@ -199,12 +226,8 @@ end
     return tupexpr(i -> :(rand(V)), N)
 end
 
-@generated function scale_tuple(tup::NTuple{N}, x) where N
-    return tupexpr(i -> :(_mul_partial(tup[$i], x)), N)
-end
-
-@generated function div_tuple_by_scalar(tup::NTuple{N}, x) where N
-    return tupexpr(i -> :(_div_partial(tup[$i], x)), N)
+@generated function scale_tuple(f::F, tup::NTuple{N}, x) where {F,N}
+    return tupexpr(i -> :(f(tup[$i], x)), N)
 end
 
 @generated function add_tuples(a::NTuple{N}, b::NTuple{N})  where N
@@ -219,8 +242,8 @@ end
     return tupexpr(i -> :(-tup[$i]), N)
 end
 
-@generated function mul_tuples(a::NTuple{N}, b::NTuple{N}, afactor, bfactor) where N
-    return tupexpr(i -> :(_mul_partial(a[$i], afactor) + _mul_partial(b[$i], bfactor)), N)
+@generated function mul_tuples(f::F, a::NTuple{N}, b::NTuple{N}, afactor, bfactor) where {F,N}
+    return tupexpr(i -> :(f(a[$i], afactor) + f(b[$i], bfactor)), N)
 end
 
 ###################

src/prelude.jl

@@ -1,4 +1,5 @@
-const NANSAFE_MODE_ENABLED = @load_preference("nansafe_mode", false)
+const NANSAFE_MODE_PREFERENCE = @load_preference("nansafe_mode", nothing) # deprecated
+const UNSAFE_MODE_ENABLED = @load_preference("unsafe_mode", NANSAFE_MODE_PREFERENCE === false)
 const DEFAULT_CHUNK_THRESHOLD = @load_preference("default_chunk_threshold", 12)
 
 const AMBIGUOUS_TYPES = (AbstractFloat, Irrational, Integer, Rational, Real, RoundingMode)

test/DerivativeTest.jl

@@ -116,13 +116,20 @@ end
 
 @testset "NaN-safe mode" begin
     x = ForwardDiff.derivative(log ∘ zero, 1.0)
-    if ForwardDiff.NANSAFE_MODE_ENABLED
+    if !ForwardDiff.UNSAFE_MODE_ENABLED
         @test iszero(x)
     else
         @test isnan(x)
     end
 end
 
+@testset "division by tiny and huge numbers" begin
+    for x0 in (1e-200, 1e-310, 1e200)
+        @test iszero(ForwardDiff.derivative(x -> 2x / x, x0))
+        @test iszero(ForwardDiff.derivative(x -> zero(x) / x, x0))
+    end
+end
+
 @testset "Givens rotations: Derivatives" begin
     # Test different branches in `LinearAlgebra.givensAlgorithm`
     for f in [randexp(), -randexp()], g in [0.0, f / 2, 2f, -f / 2, -2f], i in 1:3

test/GradientTest.jl

@@ -149,7 +149,7 @@ end
 end
 
 @testset "exponential function at base zero" begin
-    if ForwardDiff.NANSAFE_MODE_ENABLED
+    if !ForwardDiff.UNSAFE_MODE_ENABLED
         @test isequal(ForwardDiff.gradient(t -> t[1]^t[2], [0.0, -0.5]), [-Inf, -Inf])
         @test isequal(ForwardDiff.gradient(t -> t[1]^t[2], [0.0,  0.0]), [NaN, -Inf])
         @test isequal(ForwardDiff.gradient(t -> t[1]^t[2], [0.0,  0.5]), [Inf, 0.0])
@@ -214,15 +214,15 @@ end
 end
 
 @testset "gradient for exponential with NaNMath" begin
-    if ForwardDiff.NANSAFE_MODE_ENABLED
+    if !ForwardDiff.UNSAFE_MODE_ENABLED
         @test isequal(ForwardDiff.gradient(x -> NaNMath.pow(x[1],x[2]), [NaN, 1.0]), [1.0, NaN])
     else
         @test isequal(ForwardDiff.gradient(x -> NaNMath.pow(x[1],x[2]), [NaN, 1.0]), [NaN, NaN])
     end
     @test ForwardDiff.gradient(x -> NaNMath.pow(x[1], x[2]), [1.0, 1.0]) == [1.0, 0.0]
     @test isnan(ForwardDiff.gradient((x) -> NaNMath.pow(x[1], x[2]), [-1.0, 0.5])[1])
 
-    if ForwardDiff.NANSAFE_MODE_ENABLED
+    if !ForwardDiff.UNSAFE_MODE_ENABLED
         @test isequal(ForwardDiff.gradient(x -> x[1]^x[2], [NaN, 1.0]), [1.0, NaN])
     else
         @test isequal(ForwardDiff.gradient(x -> x[1]^x[2], [NaN, 1.0]), [NaN, NaN])
@@ -309,7 +309,7 @@ end
     y2 = ForwardDiff.gradient(f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{1}()))
     y3 = ForwardDiff.gradient(f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{2}()))
     for y in (y1, y2, y3)
-        if ForwardDiff.NANSAFE_MODE_ENABLED
+        if !ForwardDiff.UNSAFE_MODE_ENABLED
             @test y == [0.0, Inf]
         else
             @test isequal(y, [NaN, Inf])
@@ -323,7 +323,7 @@ end
     b2 = ForwardDiff.gradient(g, a, ForwardDiff.GradientConfig(g, a, ForwardDiff.Chunk{1}()))
     b3 = ForwardDiff.gradient(g, a, ForwardDiff.GradientConfig(g, a, ForwardDiff.Chunk{2}()))
     for b in (b1, b2, b3)
-        if ForwardDiff.NANSAFE_MODE_ENABLED
+        if !ForwardDiff.UNSAFE_MODE_ENABLED
             @test b == [Inf, -Inf]
         else
             @test isequal(b, [NaN, NaN])