Add and improve LogExpFunctions rules

Project.toml

@@ -1,6 +1,6 @@
 name = "DiffRules"
 uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
-version = "1.17.0"
+version = "1.17.1"
 
 [deps]
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
@@ -11,7 +11,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [compat]
 IrrationalConstants = "0.1.1, 0.2"
-LogExpFunctions = "0.3.2, 1"
+LogExpFunctions = "1"
 NaNMath = "0.3, 1"
 Random = "1"
 SpecialFunctions = "1.1, 2"

src/rules.jl

@@ -291,25 +291,37 @@ _abs_deriv(x) = signbit(x) ? -one(x) : one(x)
 
 # unary
 @define_diffrule LogExpFunctions.xlogx(x) = :(1 + log($x))
-@define_diffrule LogExpFunctions.logistic(x) = :(z = LogExpFunctions.logistic($x); z * (1 - z))
+# logistic(x) * logistic(-x), avoiding cancellation in `1 - logistic(x)`
+@define_diffrule LogExpFunctions.logistic(x) = :(let e = exp(-abs($x)); e / (1 + e)^2 end)
 @define_diffrule LogExpFunctions.logit(x) = :(inv($x * (1 - $x)))
-@define_diffrule LogExpFunctions.log1psq(x) = :(2 * $x / (1 + $x^2))
+# avoid overflow of `x^2`
+@define_diffrule LogExpFunctions.log1psq(x) = :(abs($x) < 1 ? 2 * $x / (1 + $x^2) : 2 / ($x + inv($x)))
 @define_diffrule LogExpFunctions.log1pexp(x) = :(LogExpFunctions.logistic($x))
-@define_diffrule LogExpFunctions.log1mexp(x) = :(-exp($x - LogExpFunctions.log1mexp($x)))
+@define_diffrule LogExpFunctions.log1mexp(x) = :(-inv(expm1(-$x)))
 @define_diffrule LogExpFunctions.log2mexp(x) = :(-exp($x - LogExpFunctions.log2mexp($x)))
-@define_diffrule LogExpFunctions.logexpm1(x) = :(exp($x - LogExpFunctions.logexpm1($x)))
+@define_diffrule LogExpFunctions.logexpm1(x) = :(-inv(expm1(-$x)))
 @define_diffrule LogExpFunctions.log1pmx(x) = :(-$x / (1 + $x))
 @define_diffrule LogExpFunctions.logmxp1(x) = :((1 - $x) / $x)
+@define_diffrule LogExpFunctions.logcosh(x) = :(tanh($x))
+@define_diffrule LogExpFunctions.logabssinh(x) = :(coth($x))
+@define_diffrule LogExpFunctions.logabstanh(x) = :(2 / sinh(2 * $x))
+@define_diffrule LogExpFunctions.loglogistic(x) = :(LogExpFunctions.logistic(-$x))
+@define_diffrule LogExpFunctions.log1mlogistic(x) = :(-LogExpFunctions.logistic($x))
+@define_diffrule LogExpFunctions.logitexp(x) = :(-inv(expm1($x)))
+@define_diffrule LogExpFunctions.logit1mexp(x) = :(inv(expm1($x)))
 
 # binary
 @define_diffrule LogExpFunctions.xlogy(x, y) =
     :(log($y)),
     :(z = $x / $y; iszero($x) && !isnan($y) ? zero(z) : z)
+# `x - y` is `NaN` if `x == y == ±Inf`
 @define_diffrule LogExpFunctions.logaddexp(x, y) =
-    :(exp($x - LogExpFunctions.logaddexp($x, $y))), :(exp($y - LogExpFunctions.logaddexp($x, $y)))
+    :(LogExpFunctions.logistic($x == $y && isinf($x) ? zero($x - $y) : $x - $y)),
+    :(LogExpFunctions.logistic($x == $y && isinf($x) ? zero($x - $y) : $y - $x))
+# `x - y` is `NaN` if `x == y == -Inf`
 @define_diffrule LogExpFunctions.logsubexp(x, y) =
-    :(z = LogExpFunctions.logsubexp($x, $y); $x > $y ? exp($x - z) : -exp($x - z)),
-    :(z = LogExpFunctions.logsubexp($x, $y); $x > $y ? -exp($y - z) : exp($y - z))
+    :(-inv(expm1(-($x == $y == -Inf ? zero($x - $y) : $x - $y)))),
+    :(-inv(expm1($x == $y == -Inf ? zero($x - $y) : $x - $y)))
 @define_diffrule LogExpFunctions.xlog1py(x, y) =
     :(log1p($y)),
     :(z = $x / (1 + $y); iszero($x) && !isnan($y) ? zero(z) : z)

test/Project.toml

@@ -11,7 +11,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 [compat]
 FiniteDifferences = "0.12"
 IrrationalConstants = "0.1.1, 0.2"
-LogExpFunctions = "0.3.2, 1"
+LogExpFunctions = "1"
 NaNMath = "0.3, 1"
 Random = "1"
 SpecialFunctions = "1.1, 2"

test/runtests.jl

@@ -27,11 +27,14 @@ non_diffeable_arg_functions = [(:Base, :rem2pi, 2), (:Base, :ldexp, 2), (:Base,
                     goo = if $(f in (:asec, :acsc, :asecd, :acscd, :acosh, :acoth))
                         # avoid singularities with finite differencing
                         rand($T) + $T(1.5)
-                    elseif $(f in (:log, :airyaix, :airyaiprimex, :logmxp1))
+                    elseif $(f in (:log, :airyaix, :airyaiprimex, :logmxp1, :logabssinh, :logabstanh))
                         # avoid singularities with finite differencing
                         rand($T) + $T(0.5)
                     elseif $(f === :log1mexp)
                         rand($T) - one($T)
+                    elseif $(f in (:logitexp, :logit1mexp))
+                        # avoid singularities with finite differencing
+                        rand($T) - $T(1.5)
                     elseif $(f in (:log2mexp, :erfinv))
                         rand($T) - $T(0.5)
                     else
@@ -40,10 +43,7 @@ non_diffeable_arg_functions = [(:Base, :rem2pi, 2), (:Base, :ldexp, 2), (:Base,
                     # We're happy with types with the correct promotion behavior, e.g.
                     # it's fine to return `1` as a derivative despite input being `Float64`.
                     @test promote_type(typeof($deriv), $T) === $T
-                    # In older versions of LogExpFunctions `log1pmx(::Float32)` and `logmxp1(::Float32)` are not defined
-                    if hasmethod($M.$f, Tuple{$T})
-                        @test $deriv ≈ finitediff($M.$f, goo) rtol=1e-2 atol=1e-3     
-                    end
+                    @test $deriv ≈ finitediff($M.$f, goo) rtol=1e-2 atol=1e-3
                     # test for 2pi functions
                     if $(f === :mod2pi)
                         goo = 4 * pi
@@ -209,6 +209,10 @@ acosh_xs(T) = (abs.(above_one(T))..., 2 * sqrt(floatmax(T)))  # x^2 overflows fo
     (:NaNMath, :acos, below_one),
     (:NaNMath, :atanh, below_one),
     (:NaNMath, :acosh, acosh_xs),
+    (:LogExpFunctions, :logistic, T -> T.((-40, 40))),
+    (:LogExpFunctions, :log1psq, T -> (-2 * sqrt(floatmax(T)), 2 * sqrt(floatmax(T)))),
+    (:LogExpFunctions, :log1mexp, T -> T.((-2e-30, -1e-12))),
+    (:LogExpFunctions, :logexpm1, T -> T.((2e-30, 6.039463507942533e-8))),
 ]
     deriv = DiffRules.diffrule(M, f, :x)
     @eval let
@@ -225,6 +229,103 @@ acosh_xs(T) = (abs.(above_one(T))..., 2 * sqrt(floatmax(T)))  # x^2 overflows fo
     end
 end
 
+# Test `Inf` arguments of `log1mexp` and `logexpm1`
+for (f, x, dx) in ((:log1mexp, -Inf, -0.0), (:logexpm1, Inf, 1.0))
+    deriv = DiffRules.diffrule(:LogExpFunctions, f, :x)
+    @eval begin
+        let
+            x = $x
+            @test $deriv === $dx
+        end
+    end
+end
+
+# Test `Inf` and `NaN` arguments of `logaddexp` (JuliaStats/LogExpFunctions.jl#61)
+derivs = DiffRules.diffrule(:LogExpFunctions, :logaddexp, :x, :y)
+for xytype in [:Float32, :Float64]
+    @eval begin
+        let
+            x = y = $xytype(-Inf)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test dx === dy === $xytype(0.5)
+
+            x = y = $xytype(Inf)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test dx === dy === $xytype(0.5)
+
+            x, y = $xytype(Inf), one($xytype)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test dx === one($xytype)
+            @test dy === zero($xytype)
+
+            x, y = $xytype(-Inf), one($xytype)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test dx === zero($xytype)
+            @test dy === one($xytype)
+
+            x, y = $xytype(NaN), one($xytype)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test isnan(dx)
+            @test isnan(dy)
+        end
+    end
+end
+
+# Test accuracy of `logaddexp` for arguments of large magnitude
+@eval begin
+    let
+        x, y = 300.0, 301.0
+        dx, dy = $(derivs[1]), $(derivs[2])
+        @test dx ≈ LogExpFunctions.logistic(big(x) - big(y)) rtol=4*eps()
+        @test dy ≈ LogExpFunctions.logistic(big(y) - big(x)) rtol=4*eps()
+    end
+end
+
+# Test equal and infinite arguments of `logsubexp`
+derivs = DiffRules.diffrule(:LogExpFunctions, :logsubexp, :x, :y)
+for xytype in [:Float32, :Float64]
+    @eval begin
+        let
+            x = y = one($xytype)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test dx === $xytype(Inf)
+            @test dy === $xytype(-Inf)
+
+            x = y = $xytype(-Inf)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test dx === $xytype(Inf)
+            @test dy === $xytype(-Inf)
+
+            x = y = $xytype(Inf)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test isnan(dx)
+            @test isnan(dy)
+
+            x, y = $xytype(Inf), one($xytype)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test dx === one($xytype)
+            @test dy === -zero($xytype)
+
+            x, y = one($xytype), $xytype(Inf)
+            dx, dy = $(derivs[1]), $(derivs[2])
+            @test dx === -zero($xytype)
+            @test dy === one($xytype)
+        end
+    end
+end
+
+# Test accuracy of `logsubexp` for nearly equal arguments
+@eval begin
+    let
+        for (x, y) in ((1.0, 1.0 + 1e-6), (1.0 + 1e-6, 1.0))
+            dx, dy = $(derivs[1]), $(derivs[2])
+            d = big(x) - big(y)
+            @test dx ≈ -inv(expm1(-d)) rtol=4*eps()
+            @test dy ≈ -inv(expm1(d)) rtol=4*eps()
+        end
+    end
+end
+
 end
 
     @testset "diffrules" begin