Merge pull request #141 from devmotion/dw/logabsgamma

Add derivative of `logabsgamma`

Project.toml

@@ -1,6 +1,6 @@
 name = "Tracker"
 uuid = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
-version = "0.2.24"
+version = "0.2.25"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/lib/real.jl

@@ -149,3 +149,13 @@ end
 
 collectmemaybe(xs::AbstractArray{>:TrackedReal}) = collect(xs)
 collectmemaybe(xs::AbstractArray{<:TrackedReal}) = collect(xs)
+
+# `logabsgamma` returns a tuple and hence its derivative is not defined in DiffRules
+SpecialFunctions.logabsgamma(x::TrackedReal) = track(SpecialFunctions.logabsgamma, x)
+@grad function SpecialFunctions.logabsgamma(x::Real)
+  data_x = data(x)
+  function logabsgamma_pullback(Δ)
+    return (SpecialFunctions.digamma(data_x) * first(Δ),)
+  end
+  return SpecialFunctions.logabsgamma(data_x), logabsgamma_pullback
+end

test/tracker.jl

@@ -2,6 +2,7 @@ using Tracker, Test, NNlib
 using Tracker: TrackedReal, gradient, gradcheck, grad, checkpoint, forwarddiff
 using NNlib: conv, ∇conv_data, depthwiseconv
 using PDMats
+using SpecialFunctions: logabsgamma
 using Printf: @sprintf
 using LinearAlgebra: diagm, dot, LowerTriangular, norm, det, logdet, logabsdet, I, Diagonal
 using Statistics: mean, std, var
@@ -504,3 +505,7 @@ end
     @test size(y) == (5, 3)
 end
 
+@testset "logabsgamma" begin
+  @test gradcheck(x -> logabsgamma(only(x))[1], rand(1))
+  @test gradcheck(x -> logabsgamma(only(x))[2], rand(1))
+end