xexpx, xexpy (#35)

Project.toml

@@ -1,7 +1,7 @@
 name = "LogExpFunctions"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 authors = ["StatsFun.jl contributors, Tamas K. Papp <[email protected]>"]
-version = "0.3.6"
+version = "0.3.7"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/index.md

@@ -10,6 +10,8 @@ LogExpFunctions supports [`InverseFunctions.inverse`](https://github.com/JuliaMa
 xlogx
 xlogy
 xlog1py
+xexpx
+xexpy
 logistic
 logit
 logcosh

src/LogExpFunctions.jl

@@ -9,7 +9,7 @@ import InverseFunctions
 import IrrationalConstants
 import LinearAlgebra
 
-export xlogx, xlogy, xlog1py, logistic, logit, log1psq, log1pexp, log1mexp, log2mexp, logexpm1,
+export xlogx, xlogy, xlog1py, xexpx, xexpy, logistic, logit, log1psq, log1pexp, log1mexp, log2mexp, logexpm1,
     softplus, invsoftplus, log1pmx, logmxp1, logaddexp, logsubexp, logsumexp, softmax,
     softmax!, logcosh
 

src/basicfuns.jl

@@ -44,6 +44,37 @@ function xlog1py(x::Number, y::Number)
     return iszero(x) && !isnan(y) ? zero(result) : result
 end
 
+"""
+$(SIGNATURES)
+
+Return `x * exp(x)` for `x > -Inf`, or zero if `x == -Inf`.
+
+```jldoctest
+julia> xexpx(-Inf)
+0.0
+```
+"""
+function xexpx(x::Real)
+    expx = exp(x)
+    return iszero(expx) ? expx : x * expx
+end
+
+"""
+$(SIGNATURES)
+
+Return `x * exp(y)` for `y > -Inf`, or zero if `y == -Inf`.
+
+```jldoctest
+julia> xexpy(1.0, -Inf)
+0.0
+```
+"""
+function xexpy(x::Real, y::Real)
+    expy = exp(y)
+    result = x * expy
+    return iszero(expy) && !isnan(x) ? zero(result) : result
+end
+
 # The following bounds are precomputed versions of the following abstract
 # function, but the implicit interface for AbstractFloat doesn't uniformly
 # enforce that all floating point types implement nextfloat and prevfloat.

src/chainrules.jl

@@ -2,6 +2,44 @@ ChainRulesCore.@scalar_rule(xlogx(x::Real), (1 + log(x),))
 ChainRulesCore.@scalar_rule(xlogy(x::Real, y::Real), (log(y), x / y,))
 ChainRulesCore.@scalar_rule(xlog1py(x::Real, y::Real), (log1p(y), x / (1 + y),))
 
+function ChainRulesCore.frule((_, Δx), ::typeof(xexpx), x::Real)
+    expx = exp(x)
+    if iszero(expx)
+        Ω = expx
+        ΔΩ = expx * Δx
+    else
+        Ω = x * expx
+        ΔΩ = (1 + x) * expx * Δx
+    end
+    return Ω, ΔΩ
+end
+
+function ChainRulesCore.rrule(::typeof(xexpx), x::Real)
+    expx = exp(x)
+    Ω = iszero(expx) ? expx : x * expx
+    function xexpx_pullback(ΔΩ)
+        Δx = iszero(expx) ? expx * ΔΩ : (1 + x) * expx * ΔΩ
+        return (ChainRulesCore.NoTangent(), Δx)
+    end
+    return Ω, xexpx_pullback
+end
+
+function ChainRulesCore.frule((_, Δx, Δy), ::typeof(xexpy), x::Real, y::Real)
+    expy = exp(y)
+    result = x * expy
+    Ω = iszero(expy) && !isnan(x) ? zero(result) : result
+    ΔΩ = expy * Δx + Ω * Δy
+    return Ω, ΔΩ
+end
+
+function ChainRulesCore.rrule(::typeof(xexpy), x::Real, y::Real)
+    expy = exp(y)
+    result = x * expy
+    Ω = iszero(expy) && !isnan(x) ? zero(result) : result
+    xexpy_pullback(ΔΩ) = (ChainRulesCore.NoTangent(), ΔΩ * expy, ΔΩ * Ω)
+    return Ω, xexpy_pullback
+end
+
 ChainRulesCore.@scalar_rule(logistic(x::Real), (Ω * (1 - Ω),))
 ChainRulesCore.@scalar_rule(logit(x::Real), (inv(x * (1 - x)),))
 ChainRulesCore.@scalar_rule(logcosh(x::Real), tanh(x))

test/basicfuns.jl

@@ -44,6 +44,37 @@
     @test iszero(xlog1py(0 + im * 0, -1 + im * Inf))
 end
 
+@testset "xexpx" begin
+    for x in (false, 0, 0.0, 0f0, -Inf, -Inf32)
+        @test (@inferred xexpx(x)) === zero(exp(x))
+    end
+    for x in (NaN16, NaN32, NaN64, Inf16, Inf32, Inf64)
+        @test (@inferred xexpx(x)) === x
+    end
+    for x in (1, true, 1.0, 1f0)
+        @test (@inferred xexpx(x)) === exp(x)
+    end
+    for a in (2, 2f0, 2.0), x in -a:a
+        @test (@inferred xexpx(x)) === x * exp(x)
+    end
+end
+
+@testset "xexpy" begin
+    for x in (0, 1, 1.0, 1f0, Inf, Inf32), y in (-Inf, -Inf32)
+        @test (@inferred xexpy(x, y)) === zero(x * exp(y))
+    end
+    for x in (0, 1, 1.0, 1f0, Inf, Inf32, -Inf, -Inf32, NaN, NaN32), nan in (NaN, NaN32)
+        @test (@inferred xexpy(x, nan)) === oftype(x * exp(nan), NaN)
+        @test (@inferred xexpy(nan, x)) === oftype(nan * exp(x), NaN)
+    end
+    for x in (2, -2f0, 2.0), y in (1, -1f0, 1.0)
+        @test (@inferred xexpy(x, y)) ≈ x * exp(y)
+    end
+    for x in (randn(), randn(Float32))
+        @test xexpy(x, x) ≈ xexpx(x)
+    end
+end
+
 @testset "logistic & logit" begin
     @test logistic(2) ≈ 1.0 / (1.0 + exp(-2.0))
     @test logistic(-750.0) === 0.0

test/chainrules.jl

@@ -13,6 +13,27 @@
         end
     end
 
+    @testset "xexpx" begin
+        # regular branch
+        test_scalar(xexpx, randn())
+        # special cases (manually since FiniteDifferences/ChainRulesTestUtils fails at -Inf)
+        @test @inferred(frule((NoTangent(), rand()), xexpx, -Inf)) === (0.0, 0.0)
+        Ω, back = @inferred(rrule(xexpx, -Inf))
+        @test Ω === 0.0
+        @test back(rand()) === (NoTangent(), 0.0)
+    end
+
+    @testset "xexpy" begin
+        # regular branch
+        test_frule(xexpy, randn(), randn())
+        test_rrule(xexpy, randn(), randn())
+        # special cases (manually since FiniteDifferences/ChainRulesTestUtils fails at -Inf)
+        @test @inferred(frule((NoTangent(), rand(), rand()), xexpy, x, -Inf)) === (0.0, 0.0)
+        Ω, back = @inferred(rrule(xexpy, x, -Inf))
+        @test Ω === 0.0
+        @test back(rand()) === (NoTangent(), 0.0, 0.0)
+    end
+
     test_frule(logit, x)
     test_rrule(logit, x)
 

test/runtests.jl

@@ -1,5 +1,6 @@
 using LogExpFunctions
 using ChainRulesTestUtils
+using ChainRulesCore
 using ChangesOfVariables
 using InverseFunctions
 using OffsetArrays