Merge pull request #25 from TuringLang/mt/zygote_multi_and_matrix

Zygote support

Author: mohamed82008

.travis.yml

@@ -8,6 +8,7 @@ julia:
   - 1.0
   - 1.1
   - 1.2
+  - 1.3
   - nightly
 
 matrix:

Project.toml

@@ -9,9 +9,11 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 Combinatorics = "0.7"
@@ -20,12 +22,12 @@ ForwardDiff = "0.10.6"
 PDMats = "0.9"
 StatsFuns = "0.8, 0.9"
 Tracker = "0.2.5"
-Zygote = "0.4.1"
+Zygote = "0.4.7"
 julia = "1"
 
 [extras]
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "FiniteDifferences"]
+test = ["Test", "FiniteDifferences"]

README.md

@@ -5,7 +5,7 @@
 [![Coverage Status](https://coveralls.io/repos/github/TuringLang/DistributionsAD.jl/badge.svg?branch=master)](https://coveralls.io/github/TuringLang/DistributionsAD.jl?branch=master)
 
 
-This package defines the necessary functions to enable automatic differentiation (AD) of the `logpdf` function from [Distributions.jl](https://github.com/JuliaStats/Distributions.jl) using the packages [Tracker.jl](https://github.com/FluxML/Tracker.jl) and [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl). The goal of this package is to make the output of `logpdf` differentiable wrt all continuous parameters of a distribution as well as the random variable in the case of continuous distributions.
+This package defines the necessary functions to enable automatic differentiation (AD) of the `logpdf` function from [Distributions.jl](https://github.com/JuliaStats/Distributions.jl) using the packages [Tracker.jl](https://github.com/FluxML/Tracker.jl), [Zygote.jl](https://github.com/FluxML/Zygote.jl) and [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl). The goal of this package is to make the output of `logpdf` differentiable wrt all continuous parameters of a distribution as well as the random variable in the case of continuous distributions.
 
 AD of `logpdf` is fully supported and tested for the following distributions wrt all combinations of continuous variables (distribution parameters and/or the random variable) and using all defined distribution constructors:
 - Univariate discrete
@@ -24,7 +24,6 @@ AD of `logpdf` is fully supported and tested for the following distributions wrt
     - `BetaPrime`
     - `Biweight`
     - `Cauchy`
-    - `Chernoff`
     - `Chi`
     - `Chisq`
     - `Cosine`
@@ -64,6 +63,8 @@ AD of `logpdf` is fully supported and tested for the following distributions wrt
     - `MvNormal`
 - Matrix-variate continuous
     - `MatrixBeta`
+    - `Wishart`
+    - `InverseWishart`
 
 # Get Involved
 

src/DistributionsAD.jl

@@ -3,13 +3,16 @@ module DistributionsAD
 using PDMats, 
       ForwardDiff, 
       Zygote, 
-      Tracker, 
       LinearAlgebra, 
       Distributions, 
       Random, 
-      Combinatorics
+      Combinatorics,
+      SpecialFunctions,
+      StatsFuns
 
-using Tracker: TrackedReal
+using Tracker: Tracker, TrackedReal, TrackedVector, TrackedMatrix, TrackedArray,
+                TrackedVecOrMat, track, data
+using ZygoteRules: ZygoteRules, pullback
 using LinearAlgebra: copytri!
 using Distributions: AbstractMvLogNormal, 
                      ContinuousMultivariateDistribution
@@ -30,10 +33,13 @@ export TuringScalMvNormal,
        TuringDiagMvNormal,
        TuringDenseMvNormal,
        TuringMvLogNormal,
-       TuringPoissonBinomial
+       TuringPoissonBinomial,
+       TuringWishart,
+       TuringInverseWishart
 
 include("common.jl")
 include("univariate.jl")
 include("multivariate.jl")
+include("matrixvariate.jl")
 
 end

src/common.jl

@@ -1,66 +1,65 @@
 ## Generic ##
 
 function Base.fill(
-    value::Tracker.TrackedReal,
+    value::TrackedReal,
     dims::Vararg{Union{Integer, AbstractUnitRange}},
 )
-    return Tracker.track(fill, value, dims...)
+    return track(fill, value, dims...)
 end
 Tracker.@grad function Base.fill(value::Real, dims...)
-    return fill(Tracker.data(value), dims...), function(Δ)
+    return fill(data(value), dims...), function(Δ)
         size(Δ) ≢  dims && error("Dimension mismatch")
         return (sum(Δ), map(_->nothing, dims)...)
     end
 end
 
 ## StatsFuns ##
 
-logsumexp(x::Tracker.TrackedArray) = Tracker.track(logsumexp, x)
-Tracker.@grad function logsumexp(x::Tracker.TrackedArray)
-    lse = logsumexp(Tracker.data(x))
-    return lse,
-          Δ->(Δ .* exp.(x .- lse),)
+logsumexp(x::TrackedArray) = track(logsumexp, x)
+Tracker.@grad function logsumexp(x::TrackedArray)
+    lse = logsumexp(data(x))
+    return lse, Δ -> (Δ .* exp.(x .- lse),)
 end
 
 ## Linear algebra ##
 
-LinearAlgebra.UpperTriangular(A::Tracker.TrackedMatrix) = Tracker.track(UpperTriangular, A)
+LinearAlgebra.UpperTriangular(A::TrackedMatrix) = track(UpperTriangular, A)
 Tracker.@grad function LinearAlgebra.UpperTriangular(A::AbstractMatrix)
-    return UpperTriangular(Tracker.data(A)), Δ->(UpperTriangular(Δ),)
+    return UpperTriangular(data(A)), Δ->(UpperTriangular(Δ),)
 end
 
-function LinearAlgebra.cholesky(A::Tracker.TrackedMatrix; check=true)
+function LinearAlgebra.cholesky(A::TrackedMatrix; check=true)
     factors_info = turing_chol(A, check)
     factors = factors_info[1]
-    info = Tracker.data(factors_info[2])
+    info = data(factors_info[2])
     return Cholesky{eltype(factors), typeof(factors)}(factors, 'U', info)
 end
 function turing_chol(A::AbstractMatrix, check)
     chol = cholesky(A, check=check)
     (chol.factors, chol.info)
 end
-turing_chol(A::Tracker.TrackedMatrix, check) = Tracker.track(turing_chol, A, check)
+turing_chol(A::TrackedMatrix, check) = track(turing_chol, A, check)
 Tracker.@grad function turing_chol(A::AbstractMatrix, check)
-    C, back = Zygote.pullback(unsafe_cholesky, Tracker.data(A), Tracker.data(check))
-    return (C.factors, C.info), Δ->back((factors=Tracker.data(Δ[1]),))
+    C, back = pullback(unsafe_cholesky, data(A), data(check))
+    return (C.factors, C.info), Δ->back((factors=data(Δ[1]),))
 end
 
 unsafe_cholesky(x, check) = cholesky(x, check=check)
-Zygote.@adjoint function unsafe_cholesky(Σ::Real, check)
+ZygoteRules.@adjoint function unsafe_cholesky(Σ::Real, check)
     C = cholesky(Σ; check=check)
     return C, function(Δ::NamedTuple)
         issuccess(C) || return (zero(Σ), nothing)
         (Δ.factors[1, 1] / (2 * C.U[1, 1]), nothing)
     end
 end
-Zygote.@adjoint function unsafe_cholesky(Σ::Diagonal, check)
+ZygoteRules.@adjoint function unsafe_cholesky(Σ::Diagonal, check)
     C = cholesky(Σ; check=check)
     return C, function(Δ::NamedTuple)
         issuccess(C) || (Diagonal(zero(diag(Δ.factors))), nothing)
         (Diagonal(diag(Δ.factors) .* inv.(2 .* C.factors.diag)), nothing)
     end
 end
-Zygote.@adjoint function unsafe_cholesky(Σ::Union{StridedMatrix, Symmetric{<:Real, <:StridedMatrix}}, check)
+ZygoteRules.@adjoint function unsafe_cholesky(Σ::Union{StridedMatrix, Symmetric{<:Real, <:StridedMatrix}}, check)
     C = cholesky(Σ; check=check)
     return C, function(Δ::NamedTuple)
         issuccess(C) || return (zero(Δ.factors), nothing)
@@ -75,38 +74,33 @@ Zygote.@adjoint function unsafe_cholesky(Σ::Union{StridedMatrix, Symmetric{<:Re
         return (UpperTriangular(Σ̄), nothing)
     end
 end
-  
+
 # Specialised logdet for cholesky to target the triangle directly.
 logdet_chol_tri(U::AbstractMatrix) = 2 * sum(log, U[diagind(U)])
-logdet_chol_tri(U::Tracker.TrackedMatrix) = Tracker.track(logdet_chol_tri, U)
+logdet_chol_tri(U::TrackedMatrix) = track(logdet_chol_tri, U)
 Tracker.@grad function logdet_chol_tri(U::AbstractMatrix)
-    U_data = Tracker.data(U)
+    U_data = data(U)
     return logdet_chol_tri(U_data), Δ->(Matrix(Diagonal(2 .* Δ ./ diag(U_data))),)
 end
 
-function LinearAlgebra.logdet(C::Cholesky{<:Tracker.TrackedReal, <:Tracker.TrackedMatrix})
+function LinearAlgebra.logdet(C::Cholesky{<:TrackedReal, <:TrackedMatrix})
     return logdet_chol_tri(C.U)
 end
 
 # Tracker's implementation of ldiv isn't good. We'll use Zygote's instead.
 zygote_ldiv(A::AbstractMatrix, B::AbstractVecOrMat) = A \ B
-function zygote_ldiv(A::Tracker.TrackedMatrix, B::Tracker.TrackedVecOrMat)
-    return Tracker.track(zygote_ldiv, A, B)
+function zygote_ldiv(A::TrackedMatrix, B::TrackedVecOrMat)
+    return track(zygote_ldiv, A, B)
 end
-function zygote_ldiv(A::Tracker.TrackedMatrix, B::AbstractVecOrMat)
-    return Tracker.track(zygote_ldiv, A, B)
+function zygote_ldiv(A::TrackedMatrix, B::AbstractVecOrMat)
+    return track(zygote_ldiv, A, B)
 end
-zygote_ldiv(A::AbstractMatrix, B::Tracker.TrackedVecOrMat) =  Tracker.track(zygote_ldiv, A, B)
+zygote_ldiv(A::AbstractMatrix, B::TrackedVecOrMat) =  track(zygote_ldiv, A, B)
 Tracker.@grad function zygote_ldiv(A, B)
-    Y, back = Zygote.pullback(\, Tracker.data(A), Tracker.data(B))
-    return Y, Δ->back(Tracker.data(Δ))
+    Y, back = pullback(\, data(A), data(B))
+    return Y, Δ->back(data(Δ))
 end
 
-function Base.:\(a::Cholesky{<:Tracker.TrackedReal, <:Tracker.TrackedArray}, b::AbstractVecOrMat)
+function Base.:\(a::Cholesky{<:TrackedReal, <:TrackedArray}, b::AbstractVecOrMat)
     return (a.U \ (a.U' \ b))
 end
-
-## PDMats ##
-
-PDMats.invquad(Σ::PDiagMat, x::Tracker.TrackedVector) = sum(abs2.(x) ./ Σ.diag)
-PDMats.invquad(Σ::PDMat, x::Tracker.TrackedVector) = sum(abs2, zygote_ldiv(Σ.chol.L, x))