Don't depend on ChainRules internals (#142)

TuringLang · Dec 11, 2020 · beddedd · beddedd · devmotion · Dec 11, 2020
1 parent ef1ffba
commit beddedd
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 27 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "DistributionsAD"
 uuid = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
-version = "0.6.12"
+version = "0.6.13"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/common.jl b/src/common.jl
@@ -4,44 +4,28 @@ function turing_chol(A::AbstractMatrix, check)
     chol = cholesky(A, check=check)
     (chol.factors, chol.info)
 end
-function ChainRules.rrule(::typeof(turing_chol), A::AbstractMatrix, check)
-    factors, info = turing_chol(A, check)
-    function turing_chol_pullback(Ȳ)
-        f̄ = Ȳ[1]
-        ∂A = ChainRules.chol_blocked_rev(f̄, factors, 25, true)
-        return (ChainRules.NO_FIELDS, ∂A, ChainRules.DoesNotExist())
-    end
-    (factors,info), turing_chol_pullback
-end
 function turing_chol_back(A::AbstractMatrix, check)
-    C, dC_pullback = rrule(turing_chol, A, check)
+    C, chol_pullback = rrule(cholesky, A, Val(false), check=check)
     function back(Δ)
-        _, dC = dC_pullback(Δ)
-        (dC, nothing)
+        Ȳ = Composite{typeof(C)}((U=Δ[1]))
+        ∂C = chol_pullback(Ȳ)[2]
+        (∂C, nothing)
     end
-    C, back
+    (C.factors,C.info), back
 end
 
 function symm_turing_chol(A::AbstractMatrix, check, uplo)
     chol = cholesky(Symmetric(A, uplo), check=check)
     (chol.factors, chol.info)
 end
-function ChainRules.rrule(::typeof(symm_turing_chol), A::AbstractMatrix, check, uplo)
-    factors, info = symm_turing_chol(A, check, uplo)
-    function symm_turing_chol_pullback(Ȳ)
-        f̄ = Ȳ[1]
-        ∂A = ChainRules.chol_blocked_rev(f̄, factors, 25, true)
-        return (ChainRules.NO_FIELDS, ∂A, ChainRules.DoesNotExist(), ChainRules.DoesNotExist())
-    end
-    return (factors,info), symm_turing_chol_pullback
-end
 function symm_turing_chol_back(A::AbstractMatrix, check, uplo)
-    C, dC_pullback = rrule(symm_turing_chol, A, check, uplo)
+    C, chol_pullback = rrule(cholesky, Symmetric(A,uplo), Val(false), check=check)
     function back(Δ)
-        _, dC = dC_pullback(Δ)
-        (dC, nothing, nothing)
+        Ȳ = Composite{typeof(C)}((U=Δ[1]))
+        ∂C = chol_pullback(Ȳ)[2]
+        (∂C, nothing, nothing)
     end
-    C, back
+    (C.factors, C.info), back
 end