diff --git a/Project.toml b/Project.toml index 2b44c925..b7f781c3 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "LogExpFunctions" uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" authors = ["StatsFun.jl contributors, Tamas K. Papp "] -version = "1.0.4" +version = "1.0.5" [deps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" diff --git a/ext/LogExpFunctionsChainRulesCoreExt.jl b/ext/LogExpFunctionsChainRulesCoreExt.jl index 28fe8fe7..b8fc2bd3 100644 --- a/ext/LogExpFunctionsChainRulesCoreExt.jl +++ b/ext/LogExpFunctionsChainRulesCoreExt.jl @@ -114,23 +114,55 @@ function ChainRulesCore.rrule(::typeof(xexpy), x::Real, y::Real) return Ω, xexpy_pullback end -ChainRulesCore.@scalar_rule(logistic(x::Real), (Ω * (1 - Ω),)) +# logistic(x) * logistic(-x), avoiding cancellation in `1 - logistic(x)` +function _∂_logistic(x::Real) + e = exp(-abs(x)) + return e / (1 + e)^2 +end +# 2x / (1 + x^2), avoiding overflow of `x^2` for large `|x|` +_∂_log1psq(x::Real) = abs(x) < 1 ? 2 * x / (1 + x^2) : 2 / (x + inv(x)) + +ChainRulesCore.@scalar_rule(logistic(x::Real), (_∂_logistic(x),)) ChainRulesCore.@scalar_rule(logit(x::Real), (inv(x * (1 - x)),)) ChainRulesCore.@scalar_rule(logcosh(x::Real), tanh(x)) ChainRulesCore.@scalar_rule(logabssinh(x::Real), coth(x)) -ChainRulesCore.@scalar_rule(logabstanh(x::Real), inv(cosh(x) * sinh(x))) -ChainRulesCore.@scalar_rule(log1psq(x::Real), (2 * x / (1 + x^2),)) +ChainRulesCore.@scalar_rule(logabstanh(x::Real), 2 / sinh(2 * x)) +ChainRulesCore.@scalar_rule(log1psq(x::Real), (_∂_log1psq(x),)) ChainRulesCore.@scalar_rule(log1pexp(x::Real), (logistic(x),)) -ChainRulesCore.@scalar_rule(log1mexp(x::Real), (-exp(x - Ω),)) +ChainRulesCore.@scalar_rule(log1mexp(x::Real), (-inv(expm1(-x)),)) ChainRulesCore.@scalar_rule(log2mexp(x::Real), (-exp(x - Ω),)) -ChainRulesCore.@scalar_rule(logexpm1(x::Real), (exp(x - Ω),)) +ChainRulesCore.@scalar_rule(logexpm1(x::Real), (-inv(expm1(-x)),)) ChainRulesCore.@scalar_rule(log1pmx(x::Real), (-x / (x + 1),)) ChainRulesCore.@scalar_rule(logmxp1(x::Real), ((1 - x) / x,)) +ChainRulesCore.@scalar_rule(loglogistic(x::Real), (logistic(-x),)) +ChainRulesCore.@scalar_rule(log1mlogistic(x::Real), (-logistic(x),)) +ChainRulesCore.@scalar_rule(logitexp(x::Real), (-inv(expm1(x)),)) +ChainRulesCore.@scalar_rule(logit1mexp(x::Real), (inv(expm1(x)),)) -ChainRulesCore.@scalar_rule(logaddexp(x::Real, y::Real), (exp(x - Ω), exp(y - Ω))) +# the one with respect to the larger argument as `1 - t` (no cancellation since +# 0 ≤ t ≤ 1/2 for `logaddexp` and t ≤ 0 for `logsubexp`). +# The handling of `x == y` mirrors the primal implementations. +# (0 ≤ t ≤ 1/2 for `logaddexp` and t ≤ 0 for `logsubexp`, so `1 - t` is accurate). +# `x - y` is `NaN` if `x == y == ±Inf`. +function _∂_logaddexp(x::Real, y::Real) + Δ = x == y && isinf(x) ? zero(x - y) : abs(x - y) + t = logistic(-Δ) + return x >= y ? (1 - t, t) : (t, 1 - t) +end +function _∂_logsubexp(x::Real, y::Real) + Δ = x == y == -Inf ? zero(x - y) : abs(x - y) + t = -inv(expm1(Δ)) + return x >= y ? (1 - t, t) : (t, 1 - t) +end +ChainRulesCore.@scalar_rule( + logaddexp(x::Real, y::Real), + @setup((∂x, ∂y) = _∂_logaddexp(x, y)), + (∂x, ∂y), +) ChainRulesCore.@scalar_rule( logsubexp(x::Real, y::Real), - (x > y ? exp(x - Ω) : -exp(x - Ω), x > y ? -exp(y - Ω) : exp(y - Ω)), + @setup((∂x, ∂y) = _∂_logsubexp(x, y)), + (∂x, ∂y), ) function ChainRulesCore.frule((_, Δx), ::typeof(logsumexp), x::AbstractArray{<:Real}; dims=:) @@ -185,7 +217,7 @@ function ChainRulesCore.rrule(::typeof(softmax), x::AbstractArray{<:Real}; dims= return Ω, softmax_pullback end -ChainRulesCore.@scalar_rule(cloglog(x), (-inv((1 - x) * log1p(-x)),)) -ChainRulesCore.@scalar_rule(cexpexp(x), (-xexpx(-exp(x)),)) +ChainRulesCore.@scalar_rule(cloglog(x::Real), (-inv((1 - x) * log1p(-x)),)) +ChainRulesCore.@scalar_rule(cexpexp(x::Real), (-xexpx(-exp(x)),)) end # module diff --git a/test/chainrules.jl b/test/chainrules.jl index c19ae936..101d8855 100644 --- a/test/chainrules.jl +++ b/test/chainrules.jl @@ -48,6 +48,11 @@ test_frule(log1psq, x) test_rrule(log1psq, x) end + # no overflow of `x^2` + for x in (-1e200, 1e200) + @test last(frule((NoTangent(), 1.0), log1psq, x)) ≈ 2 / x rtol=4*eps() + @test last(last(rrule(log1psq, x))(1.0)) ≈ 2 / x rtol=4*eps() + end # test all `Float64` and `Float32` branches of `logistic` for x in (-821.4, -23.5, 12.3, 41.2) @@ -58,6 +63,13 @@ test_frule(logistic, x; rtol=1f-3, atol=1f-3) test_rrule(logistic, x; rtol=1f-3, atol=1f-3) end + # derivative is accurate for large `|x|` in a relative sense (no cancellation) + for x in (40.0, -40.0, 30f0, -30f0) + e = exp(-abs(big(x))) + ∂x = e / (1 + e)^2 + @test last(frule((NoTangent(), one(x)), logistic, x)) ≈ ∂x rtol=4*eps(typeof(x)) + @test last(last(rrule(logistic, x))(one(x))) ≈ ∂x rtol=4*eps(typeof(x)) + end for x in (-randexp(), randexp()) test_frule(logcosh, x) @@ -94,6 +106,25 @@ test_frule(logexpm1, x; rtol=1f-3, atol=1f-3) test_rrule(logexpm1, x; rtol=1f-3, atol=1f-3) end + # special case (manually since FiniteDifferences/ChainRulesTestUtils fails at Inf) + @test frule((NoTangent(), 1.0), logexpm1, Inf) === (Inf, 1.0) + @test last(last(rrule(logexpm1, Inf))(1.0)) === 1.0 + + # accuracy for small `|x|` (where `Ω` is large in magnitude) + for (f, x) in ((log1mexp, -1e-12), (logexpm1, 6.039463507942533e-8)) + ∂x = -inv(expm1(-big(x))) + @test last(frule((NoTangent(), 1.0), f, x)) ≈ ∂x rtol=4*eps() + @test last(last(rrule(f, x))(1.0)) ≈ ∂x rtol=4*eps() + end + + for x in (-randexp(), randexp()) + test_scalar(loglogistic, x) + test_scalar(log1mlogistic, x) + end + for x in (-10.2, -3.3, -0.3) + test_scalar(logitexp, x) + test_scalar(logit1mexp, x) + end test_scalar(log1pmx, rand()) @@ -106,6 +137,55 @@ test_frule(logsubexp, x, y) test_rrule(logsubexp, x, y) end + # nearly equal arguments (manually since finite differences are inaccurate close to + # the singularity at `x == y`) + for (x, y) in ((1.0, 1.0 + 1e-6), (1.0 + 1e-6, 1.0)) + d = big(x) - big(y) + ∂x, ∂y = -inv(expm1(-d)), -inv(expm1(d)) + @test last(frule((NoTangent(), 1.0, 0.0), logsubexp, x, y)) ≈ ∂x rtol=4*eps() + @test last(frule((NoTangent(), 0.0, 1.0), logsubexp, x, y)) ≈ ∂y rtol=4*eps() + _, back = rrule(logsubexp, x, y) + _, x̄, ȳ = back(1.0) + @test x̄ ≈ ∂x rtol=4*eps() + @test ȳ ≈ ∂y rtol=4*eps() + end + # accuracy for arguments of large magnitude + let (x, y) = (300.0, 301.0) + d = big(x) - big(y) + ∂x, ∂y = logistic(d), logistic(-d) + _, back = rrule(logaddexp, x, y) + _, x̄, ȳ = back(1.0) + @test x̄ ≈ ∂x rtol=4*eps() + @test ȳ ≈ ∂y rtol=4*eps() + end + # special cases (manually since FiniteDifferences/ChainRulesTestUtils fails at Inf) + for (f, ∂x, ∂y) in ((logaddexp, 1.0, 0.0), (logsubexp, 1.0, -0.0)) + Ω, back = rrule(f, Inf, 1.0) + @test Ω === Inf + @test back(1.0) === (NoTangent(), ∂x, ∂y) + Ω, back = rrule(f, 1.0, Inf) + @test Ω === Inf + @test back(1.0) === (NoTangent(), ∂y, ∂x) + @test last(frule((NoTangent(), 1.0, 0.0), f, Inf, 1.0)) === ∂x + end + # equal and `NaN` arguments (`x == y == ±Inf`: #61) + for T in (Float32, Float64) + for x in (T(-Inf), one(T), T(Inf)) + _, back = rrule(logaddexp, x, x) + @test back(one(T)) === (NoTangent(), T(0.5), T(0.5)) + end + # limits as `x` approaches `y` from above + for x in (T(-Inf), one(T)) + _, back = rrule(logsubexp, x, x) + @test back(one(T)) === (NoTangent(), T(Inf), T(-Inf)) + end + _, back = rrule(logsubexp, T(Inf), T(Inf)) + @test all(isnan, back(one(T))[2:3]) + for f in (logaddexp, logsubexp) + _, back = rrule(f, T(NaN), one(T)) + @test all(isnan, back(one(T))[2:3]) + end + end for x in (randn(10), randn(10, 8)), dims in (:, 1, 1:2, 2) dims isa Colon || all(d <= ndims(x) for d in dims) || continue