Avoid catastrophic cancellation in tanh and inverse trig/hyperbolic derivatives

Project.toml

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

src/rulesets/Base/base.jl

@@ -168,22 +168,34 @@ end
 
 # Can't multiply though sqrt in acosh because of negative complex case for x
 @scalar_rule acosh(x) inv(sqrt(x - 1) * sqrt(x + 1))
-@scalar_rule acoth(x) inv(1 - x ^ 2)
+# `(1 - x) * (1 + x)` avoids cancellation in `1 - x ^ 2` for `|x|` close to 1
+# (Goldberg (1991), "What Every Computer Scientist Should Know About Floating-Point Arithmetic")
+@scalar_rule acoth(x) inv((1 - x) * (1 + x))
 @scalar_rule acsch(x) -(inv(x ^ 2 * sqrt(1 + x ^ -2)))
 @scalar_rule acsch(x::Real) -(inv(abs(x) * sqrt(1 + x ^ 2)))
-@scalar_rule asech(x) -(inv(x * sqrt(1 - x ^ 2)))
+@scalar_rule asech(x) -inv(x * sqrt((1 - x) * (1 + x)))
 @scalar_rule asinh(x) inv(sqrt(x ^ 2 + 1))
-@scalar_rule atanh(x) inv(1 - x ^ 2)
+# if `x ^ 2` overflows, `1 / |x|` is exact up to rounding
+@scalar_rule asinh(x::Real) (x2=x ^ 2; isfinite(x2) ? inv(sqrt(x2 + 1)) : inv(abs(x)))
+# `(1 + ix) * (1 - ix)` avoids cancellation in `1 + x ^ 2` near `x = ±i`;
+# `ix` is constructed explicitly since `im * x` loses the sign of zero
+@scalar_rule asinh(x::Complex) @setup(ix=Complex(-imag(x), real(x))) inv(
+    sqrt((1 + ix) * (1 - ix))
+)
+@scalar_rule atanh(x) inv((1 - x) * (1 + x))
 
 
-@scalar_rule acosd(x) -inv(deg2rad(sqrt(1 - x ^ 2)))
+@scalar_rule acosd(x) -inv(deg2rad(sqrt((1 - x) * (1 + x))))
 @scalar_rule acotd(x) -inv(deg2rad(1 + x ^ 2))
+# `(1 + im * x) * (1 - im * x)` avoids cancellation in `1 + x ^ 2` near `x = ±i`
+@scalar_rule acotd(x::Complex) -inv(deg2rad((1 + im * x) * (1 - im * x)))
 @scalar_rule acscd(x) -inv(deg2rad(x^2 * sqrt(1 - x ^ -2)))
-@scalar_rule acscd(x::Real) -inv(deg2rad(abs(x) * sqrt(x ^ 2 - 1)))
+@scalar_rule acscd(x::Real) -inv(deg2rad(abs(x) * sqrt((x - 1) * (x + 1))))
 @scalar_rule asecd(x) inv(deg2rad(x ^ 2 * sqrt(1 - x ^ -2)))
-@scalar_rule asecd(x::Real) inv(deg2rad(abs(x) * sqrt(x ^ 2 - 1)))
-@scalar_rule asind(x) inv(deg2rad(sqrt(1 - x ^ 2)))
+@scalar_rule asecd(x::Real) inv(deg2rad(abs(x) * sqrt((x - 1) * (x + 1))))
+@scalar_rule asind(x) inv(deg2rad(sqrt((1 - x) * (1 + x))))
 @scalar_rule atand(x) inv(deg2rad(1 + x ^ 2))
+@scalar_rule atand(x::Complex) inv(deg2rad((1 + im * x) * (1 - im * x)))
 
 @scalar_rule cot(x) -((1 + Ω ^ 2))
 @scalar_rule coth(x) -(csch(x) ^ 2)
@@ -196,10 +208,11 @@ end
 @scalar_rule sech(x) -(tanh(x)) * Ω
 
 @scalar_rule acot(x) -(inv(1 + x ^ 2))
+@scalar_rule acot(x::Complex) -inv((1 + im * x) * (1 - im * x))
 @scalar_rule acsc(x) -(inv(x ^ 2 * sqrt(1 - x ^ -2)))
-@scalar_rule acsc(x::Real) -(inv(abs(x) * sqrt(x ^ 2 - 1)))
+@scalar_rule acsc(x::Real) -inv(abs(x) * sqrt((x - 1) * (x + 1)))
 @scalar_rule asec(x) inv(x ^ 2 * sqrt(1 - x ^ -2))
-@scalar_rule asec(x::Real) inv(abs(x) * sqrt(x ^ 2 - 1))
+@scalar_rule asec(x::Real) inv(abs(x) * sqrt((x - 1) * (x + 1)))
 
 @scalar_rule cosd(x) -deg2rad(sind(x))
 @scalar_rule cospi(x) -π * sinpi(x)

src/rulesets/Base/fastmath_able.jl

@@ -38,16 +38,23 @@ let
         # Trig-Hyperbolic
         @scalar_rule cosh(x) sinh(x)
         @scalar_rule sinh(x) cosh(x)
-        @scalar_rule tanh(x) 1 - Ω ^ 2
+        # `1 - Ω ^ 2` cancels catastrophically for large `|x|` (DiffRules.jl#102)
+        @scalar_rule tanh(x) abs(real(x)) < 1 ? 1 - Ω ^ 2 : sech(x) ^ 2
 
         # Trig- Inverses
-        @scalar_rule acos(x) -(inv(sqrt(1 - x ^ 2)))
-        @scalar_rule asin(x) inv(sqrt(1 - x ^ 2))
+        # `(1 - x) * (1 + x)` avoids cancellation in `1 - x ^ 2` for `|x|` close to 1 (see `acoth` in base.jl)
+        @scalar_rule acos(x) -inv(sqrt((1 - x) * (1 + x)))
+        @scalar_rule asin(x) inv(sqrt((1 - x) * (1 + x)))
         @scalar_rule atan(x) inv(1 + x ^ 2)
+        # `(1 + im * x) * (1 - im * x)` avoids cancellation in `1 + x ^ 2` near `x = ±i`
+        @scalar_rule atan(x::Complex) inv((1 + im * x) * (1 - im * x))
 
         # Trig-Multivariate
-        @scalar_rule atan(y, x) @setup(u = x ^ 2 + y ^ 2) (x / u, -y / u)
-        @scalar_rule sincos(x) @setup((sinx, cosx) = Ω) cosx -sinx
+        # scaling by `max(|x|, |y|)` avoids overflow and underflow of `x ^ 2 + y ^ 2`
+        @scalar_rule atan(y, x) @setup(
+            s=max(abs(x), abs(y)), xs=x / s, ys=y / s, v=xs ^ 2 + ys ^ 2
+        ) (xs / v / s, -ys / v / s)
+        @scalar_rule sincos(x) @setup((sinx, cosx)=Ω) cosx -sinx
 
         # exponents
         @scalar_rule cbrt(x) inv(3 * Ω ^ 2)

test/rulesets/Base/base.jl

@@ -67,6 +67,54 @@ end
             test_scalar(acoth, x + 1)
         end
 
+        # finite differencing is not accurate enough to detect cancellation near |x| = 1
+        @testset "Inverse accuracy near |x| = 1" begin
+            setprecision(BigFloat, 256) do
+                for (f, df, xs) in (
+                    (atanh, x -> inv(1 - x^2), (1 - 1e-8, -1 + 1e-8, ComplexF64(1, 1e-9))),
+                    (acoth, x -> inv(1 - x^2), (1 + 1e-8, -1 - 1e-8, ComplexF64(1, 1e-9))),
+                    (asech, x -> -inv(x * sqrt(1 - x^2)), (1 - 1e-8, ComplexF64(1, 1e-9))),
+                    (asind, x -> inv(deg2rad(sqrt(1 - x^2))), (1 - 1e-8, -1 + 1e-8)),
+                    (acosd, x -> -inv(deg2rad(sqrt(1 - x^2))), (1 - 1e-8, -1 + 1e-8)),
+                    (
+                        asecd,
+                        x -> inv(deg2rad(abs(x) * sqrt(x^2 - 1))),
+                        (1 + 1e-8, -1 - 1e-8),
+                    ),
+                    (
+                        acscd,
+                        x -> -inv(deg2rad(abs(x) * sqrt(x^2 - 1))),
+                        (1 + 1e-8, -1 - 1e-8),
+                    ),
+                    (asec, x -> inv(abs(x) * sqrt(x^2 - 1)), (1 + 1e-8, -1 - 1e-8)),
+                    (acsc, x -> -inv(abs(x) * sqrt(x^2 - 1)), (1 + 1e-8, -1 - 1e-8)),
+                )
+                    @testset "$f($x)" for x in xs
+                        _, ẏ = frule((NoTangent(), one(x)), f, x)
+                        @test ẏ ≈ df(big(x)) rtol = 4eps()
+                    end
+                end
+            end
+        end
+
+        # finite differencing is not accurate enough to detect cancellation near x = ±i
+        @testset "Inverse accuracy near x = ±i and for large |x|" begin
+            setprecision(BigFloat, 256) do
+                zs = (ComplexF64(1e-9, 1 - 1e-8), ComplexF64(-1e-9, -1 + 1e-8))
+                for (f, df, xs) in (
+                    (asinh, x -> inv(sqrt(1 + x^2)), (zs..., 1e200, -1e200)),
+                    (acot, x -> -inv(1 + x^2), zs),
+                    (atand, x -> inv(deg2rad(1 + x^2)), zs),
+                    (acotd, x -> -inv(deg2rad(1 + x^2)), zs),
+                )
+                    @testset "$f($x)" for x in xs
+                        _, ẏ = frule((NoTangent(), one(x)), f, x)
+                        @test ẏ ≈ df(big(x)) rtol = 4eps()
+                    end
+                end
+            end
+        end
+
         @testset "Inverse degrees" for x = (0.5, Complex(0.5, 0.25))
             test_scalar(asind, x)
             test_scalar(acosd, x)

test/rulesets/Base/fastmath_able.jl

@@ -59,6 +59,46 @@ const FASTABLE_AST = quote
             test_scalar(acos, x)
             test_scalar(atan, x)
         end
+        # finite differencing is not accurate enough to detect cancellation
+        @testset "Accuracy" begin
+            setprecision(BigFloat, 256) do
+                for (f, df, xs) in (
+                    (
+                        tanh,
+                        x -> sech(x)^2,
+                        (
+                            0.5,
+                            4.0,
+                            14.0,
+                            25.0,
+                            -25.0,
+                            ComplexF64(14, 0.3),
+                            ComplexF64(30, 1),
+                        ),
+                    ),
+                    (
+                        asin,
+                        x -> inv(sqrt(1 - x^2)),
+                        (1 - 1e-8, -1 + 1e-8, ComplexF64(1, 1e-9)),
+                    ),
+                    (
+                        acos,
+                        x -> -inv(sqrt(1 - x^2)),
+                        (1 - 1e-8, -1 + 1e-8, ComplexF64(1, 1e-9)),
+                    ),
+                    (
+                        atan,
+                        x -> inv(1 + x^2),
+                        (ComplexF64(1e-9, 1 - 1e-8), ComplexF64(-1e-9, -1 + 1e-8)),
+                    ),
+                )
+                    @testset "$f($x)" for x in xs
+                        _, ẏ = frule((NoTangent(), one(x)), f, x)
+                        @test ẏ ≈ df(big(x)) rtol = 4eps()
+                    end
+                end
+            end
+        end
         @testset "Multivariate" begin
             @testset "sincos(x::$T)" for T in (Float64, ComplexF64)
                 Δz = Tangent{Tuple{T,T}}(randn(T), randn(T))
@@ -299,3 +339,17 @@ end
 @eval @testset "fastmath_able FastMath functions" begin
     $(Base.FastMath.make_fastmath(FASTABLE_AST))
 end
+
+@testset "atan(y, x) for large and small arguments" begin
+    setprecision(BigFloat, 256) do
+        @testset "$f($y, $x)" for f in (atan, Base.FastMath.atan_fast),
+            (y, x) in ((1e-200, 1e-200), (1e200, 3e200), (-1e-170, 3e-170))
+
+            _, pb = rrule(f, y, x)
+            _, ȳ, x̄ = pb(1.0)
+            u = big(x)^2 + big(y)^2
+            @test ȳ ≈ big(x) / u rtol = 4eps()
+            @test x̄ ≈ -big(y) / u rtol = 4eps()
+        end
+    end
+end