Add ChainRules definitions (#25)

Author: devmotion

Project.toml

@@ -1,21 +1,25 @@
 name = "LogExpFunctions"
 uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 authors = ["StatsFun.jl contributors, Tamas K. Papp <[email protected]>"]
-version = "0.3.0"
+version = "0.3.1"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [compat]
+ChainRulesCore = "1"
 DocStringExtensions = "0.8"
 IrrationalConstants = "0.1"
 julia = "1"
 
 [extras]
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["OffsetArrays", "Test"]
+test = ["ChainRulesTestUtils", "OffsetArrays", "Random", "Test"]

src/LogExpFunctions.jl

@@ -3,6 +3,7 @@ module LogExpFunctions
 using DocStringExtensions: SIGNATURES
 using Base: Math.@horner
 
+import ChainRulesCore
 import IrrationalConstants
 import LinearAlgebra
 
@@ -12,5 +13,6 @@ export xlogx, xlogy, logistic, logit, log1psq, log1pexp, log1mexp, log2mexp, log
 
 include("basicfuns.jl")
 include("logsumexp.jl")
+include("chainrules.jl")
 
 end # module

src/chainrules.jl

@@ -0,0 +1,59 @@
+ChainRulesCore.@scalar_rule(xlogx(x::Real), (1 + log(x),))
+ChainRulesCore.@scalar_rule(xlogy(x::Real, y::Real), (log(y), x / y,))
+
+ChainRulesCore.@scalar_rule(logistic(x::Real), (Ω * (1 - Ω),))
+ChainRulesCore.@scalar_rule(logit(x::Real), (inv(x * (1 - x)),))
+ChainRulesCore.@scalar_rule(log1psq(x::Real), (2 * x / (1 + x^2),))
+ChainRulesCore.@scalar_rule(log1pexp(x::Real), (logistic(x),))
+ChainRulesCore.@scalar_rule(log1mexp(x::Real), (-exp(x - Ω),))
+ChainRulesCore.@scalar_rule(log2mexp(x::Real), (-exp(x - Ω),))
+ChainRulesCore.@scalar_rule(logexpm1(x::Real), (exp(x - Ω),))
+
+ChainRulesCore.@scalar_rule(logaddexp(x::Real, y::Real), (exp(x - Ω), exp(y - Ω)))
+ChainRulesCore.@scalar_rule(
+    logsubexp(x::Real, y::Real),
+    (x > y ? exp(x - Ω) : -exp(x - Ω), x > y ? -exp(y - Ω) : exp(y - Ω)),
+)
+
+function ChainRulesCore.frule((_, Δx), ::typeof(logsumexp), x::AbstractArray{<:Real}; dims=:)
+    Ω = logsumexp(x; dims=dims)
+    ΔΩ = sum(exp.(x .- Ω) .* Δx; dims=dims)
+    return Ω, ΔΩ
+end
+function ChainRulesCore.rrule(::typeof(logsumexp), x::AbstractArray{<:Real}; dims=:)
+    Ω = logsumexp(x; dims=dims)
+    project_x = ChainRulesCore.ProjectTo(x)
+    function logsumexp_pullback(Ω̄)
+        x̄ = ChainRulesCore.InplaceableThunk(
+            Δ -> Δ .+= Ω̄ .* exp.(x .- Ω),
+            ChainRulesCore.@thunk(project_x(Ω̄ .* exp.(x .- Ω))),
+        )
+        return ChainRulesCore.NoTangent(), x̄
+    end
+    return Ω, logsumexp_pullback
+end
+
+function ChainRulesCore.frule(
+    (_, _, Δx), ::typeof(softmax!), r::AbstractArray{<:Real}, x::AbstractArray{<:Real},
+)
+    softmax!(r, x)
+    _Δx = reshape(Δx, size(r))
+    Δr = r .* (_Δx .- LinearAlgebra.dot(r, _Δx))
+    return r, Δr
+end
+function ChainRulesCore.rrule(
+    ::typeof(softmax!), r::AbstractArray{<:Real}, x::AbstractArray{<:Real},
+)
+    softmax!(r, x)
+    project_x = ChainRulesCore.ProjectTo(x)
+    rcopy = copy(reshape(r, size(x)))
+    function softmax!_pullback(r̄)
+        _r̄ = reshape(r̄, size(rcopy))
+        x̄ = ChainRulesCore.InplaceableThunk(
+            Δ -> Δ .+= rcopy .* (_r̄ .- LinearAlgebra.dot(rcopy, _r̄)),
+            ChainRulesCore.@thunk(project_x(rcopy .* (_r̄ .- LinearAlgebra.dot(rcopy, _r̄)))),
+        )
+        return ChainRulesCore.NoTangent(), ChainRulesCore.ZeroTangent(), x̄
+    end
+    return r, softmax!_pullback
+end

test/chainrules.jl

@@ -0,0 +1,79 @@
+@testset "chainrules.jl" begin
+    x = rand()
+    test_frule(xlogx, x)
+    test_rrule(xlogx, x)
+    for x in (-x, 0.0, x)
+        y = rand()
+        test_frule(xlogy, x, y)
+        test_rrule(xlogy, x, y)
+    end
+
+    test_frule(logit, x)
+    test_rrule(logit, x)
+
+    for x in (-randexp(), randexp())
+        test_frule(log1psq, x)
+        test_rrule(log1psq, x)
+    end
+
+    # test all `Float64` and `Float32` branches of `logistic`
+    for x in (-821.4, -23.5, 12.3, 41.2)
+        test_frule(logistic, x)
+        test_rrule(logistic, x)
+    end
+    for x in (-123.2f0, -21.4f0, 8.3f0, 21.5f0)
+        test_frule(logistic, x; rtol=1f-3, atol=1f-3)
+        test_rrule(logistic, x; rtol=1f-3, atol=1f-3)
+    end
+
+    # test all branches of `log1pexp`
+    for x in (-20.9, 15.4, 41.5)
+        test_frule(log1pexp, x)
+        test_rrule(log1pexp, x)
+    end
+    for x in (8.3f0, 12.5f0, 21.2f0)
+        test_frule(log1pexp, x; rtol=1f-3, atol=1f-3)
+        test_rrule(log1pexp, x; rtol=1f-3, atol=1f-3)
+    end
+
+    for x in (-10.2, -3.3, -0.3)
+        test_frule(log1mexp, x)
+        test_rrule(log1mexp, x)
+    end
+
+    for x in (-10.2, -3.3, -0.3, 0.5)
+        test_frule(log2mexp, x)
+        test_rrule(log2mexp, x)
+    end
+
+    # test all branches of `logexpm1`
+    for x in (5.2, 21.4, 41.5)
+        test_frule(logexpm1, x)
+        test_rrule(logexpm1, x)
+    end
+    for x in (4.3f0, 12.5f0, 21.2f0)
+        test_frule(logexpm1, x; rtol=1f-3, atol=1f-3)
+        test_rrule(logexpm1, x; rtol=1f-3, atol=1f-3)
+    end
+
+    for x in (-randexp(), randexp()), y in (-randexp(), randexp())
+        test_frule(logaddexp, x, y)
+        test_rrule(logaddexp, x, y)
+
+        test_frule(logsubexp, x, y)
+        test_rrule(logsubexp, x, y)
+    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
+        test_frule(logsumexp, x; fkwargs=(dims=dims,))
+        test_rrule(logsumexp, x; fkwargs=(dims=dims,))
+    end
+
+    for x in (randn(10), randn(10, 8))
+        for r in (similar(x), similar(x, 1, size(x)...))
+            test_frule(softmax!, r, x)
+            test_rrule(softmax!, r, x)
+        end
+    end
+end

test/runtests.jl

@@ -1,5 +1,11 @@
 using LogExpFunctions
+using ChainRulesTestUtils
 using OffsetArrays
+
+using Random
 using Test
 
+Random.seed!(1234)
+
 include("basicfuns.jl")
+include("chainrules.jl")