Numerical improvements to correlation bijectors (#313)

* Rename VecCholeskyBijector to VecCorrCholeskyBijector * Compute corr logdetjac during transform * Enforce one-based indexing * Add with_logabsdet_jacobian for correlation transforms * Add rrule for non-mutating ADs * Update ChainRules to use manual rrule * Update Tracker to use manual rrule * Remove rrule for ReverseDiff `@grad_from_chainrules` can't handle multi-output functions, see JuliaDiff/ReverseDiff.jl#221. In this case it can AD through the primal just fine. * Add module * Make CorrBijector more numerically stable Also use consistent notation with inverse transform * Increment patch number * Revert "Rename VecCholeskyBijector to VecCorrCholeskyBijector" This reverts commit bd6ff3d. * Update src/bijectors/corr.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Tor Erlend Fjelde <[email protected]> * Apply suggestions from code review * Work around issues with Tracker * import `stack` from Compat.jl (#314) * import `stack` in tests too * disable certain tests for ProductBijector on Julia versions with older `eachslice` impls --------- Co-authored-by: Tor Erlend Fjelde <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
TuringLang · Jun 5, 2024 · fd53666 · fd53666 · torfjelde · Jun 5, 2024
1 parent 1070b73
commit fd53666
Show file tree

Hide file tree

Showing 10 changed files with 194 additions and 202 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Bijectors"
 uuid = "76274a88-744f-5084-9051-94815aaf08c4"
-version = "0.13.12"
+version = "0.13.13"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
@@ -43,7 +43,7 @@ ArgCheck = "1, 2"
 ChainRules = "1"
 ChainRulesCore = "0.10.11, 1"
 ChangesOfVariables = "0.1"
-Compat = "3, 4"
+Compat = "3.46, 4.2"
 Distributions = "0.25.33"
 ForwardDiff = "0.10"
 DistributionsAD = "0.6"

diff --git a/ext/BijectorsReverseDiffExt.jl b/ext/BijectorsReverseDiffExt.jl
@@ -268,7 +268,6 @@ end
 @grad_from_chainrules _link_chol_lkj(x::TrackedMatrix)
 @grad_from_chainrules _link_chol_lkj_from_upper(x::TrackedMatrix)
 @grad_from_chainrules _link_chol_lkj_from_lower(x::TrackedMatrix)
-@grad_from_chainrules _inv_link_chol_lkj(x::TrackedVector)
 
 cholesky_lower(X::TrackedMatrix) = track(cholesky_lower, X)
 @grad function cholesky_lower(X_tracked::TrackedMatrix)

diff --git a/ext/BijectorsTrackerExt.jl b/ext/BijectorsTrackerExt.jl
@@ -338,106 +338,16 @@ Bijectors.lower_triangular(A::TrackedMatrix) = track(Bijectors.lower_triangular,
 end
 
 Bijectors._inv_link_chol_lkj(y::TrackedVector) = track(Bijectors._inv_link_chol_lkj, y)
-@grad function Bijectors._inv_link_chol_lkj(y_tracked::TrackedVector)
-    y = data(y_tracked)
-    K = _triu1_dim_from_length(length(y))
-
-    W = similar(y, K, K)
-
-    z_vec = similar(y)
-    tmp_vec = similar(y)
-
-    idx = 1
-    @inbounds for j in 1:K
-        W[1, j] = 1
-        for i in 2:j
-            z = tanh(y[idx])
-            tmp = W[i - 1, j]
-
-            z_vec[idx] = z
-            tmp_vec[idx] = tmp
-            idx += 1
-
-            W[i - 1, j] = z * tmp
-            W[i, j] = tmp * sqrt(1 - z^2)
-        end
-        for i in (j + 1):K
-            W[i, j] = 0
-        end
-    end
-
-    function pullback_inv_link_chol_lkj(ΔW)
-        LinearAlgebra.checksquare(ΔW)
-
-        Δy = zero(y)
-
-        @inbounds for j in 1:K
-            idx_up_to_prev_column = ((j - 1) * (j - 2) ÷ 2)
-            Δtmp = ΔW[j, j]
-            for i in j:-1:2
-                idx = idx_up_to_prev_column + i - 1
-                Δz =
-                    ΔW[i - 1, j] * tmp_vec[idx] -
-                    Δtmp * tmp_vec[idx] / sqrt(1 - z_vec[idx]^2) * z_vec[idx]
-                Δy[idx] = Δz / cosh(y[idx])^2
-                Δtmp = ΔW[i - 1, j] * z_vec[idx] + Δtmp * sqrt(1 - z_vec[idx]^2)
-            end
-        end
-
-        return (Δy,)
-    end
-
-    return W, pullback_inv_link_chol_lkj
-end
-
 Bijectors._inv_link_chol_lkj(y::TrackedMatrix) = track(Bijectors._inv_link_chol_lkj, y)
-@grad function Bijectors._inv_link_chol_lkj(y_tracked::TrackedMatrix)
+@grad function Bijectors._inv_link_chol_lkj(y_tracked::Union{TrackedVector,TrackedMatrix})
     y = data(y_tracked)
+    W_logJ, back = Bijectors._inv_link_chol_lkj_rrule(y)
 
-    K = LinearAlgebra.checksquare(y)
-
-    w = similar(y)
-
-    z_mat = similar(y) # cache for adjoint
-    tmp_mat = similar(y)
-
-    @inbounds for j in 1:K
-        w[1, j] = 1
-        for i in 2:j
-            z = tanh(y[i - 1, j])
-            tmp = w[i - 1, j]
-
-            z_mat[i, j] = z
-            tmp_mat[i, j] = tmp
-
-            w[i - 1, j] = z * tmp
-            w[i, j] = tmp * sqrt(1 - z^2)
-        end
-        for i in (j + 1):K
-            w[i, j] = 0
-        end
-    end
-
-    function pullback_inv_link_chol_lkj(Δw)
-        LinearAlgebra.checksquare(Δw)
-
-        Δy = zero(y)
-
-        @inbounds for j in 1:K
-            Δtmp = Δw[j, j]
-            for i in j:-1:2
-                Δz =
-                    Δw[i - 1, j] * tmp_mat[i, j] -
-                    Δtmp * tmp_mat[i, j] / sqrt(1 - z_mat[i, j]^2) * z_mat[i, j]
-                Δy[i - 1, j] = Δz / cosh(y[i - 1, j])^2
-                Δtmp = Δw[i - 1, j] * z_mat[i, j] + Δtmp * sqrt(1 - z_mat[i, j]^2)
-            end
-        end
-
-        return (Δy,)
+    function pullback_inv_link_chol_lkj(ΔW_ΔlogJ)
+        return (back(ΔW_ΔlogJ),)
     end
 
-    return w, pullback_inv_link_chol_lkj
+    return W_logJ, pullback_inv_link_chol_lkj
 end
 
 Bijectors._link_chol_lkj(w::TrackedMatrix) = track(Bijectors._link_chol_lkj, w)

diff --git a/src/Bijectors.jl b/src/Bijectors.jl
@@ -82,6 +82,10 @@ if VERSION < v"1.1"
     using Compat: eachcol
 end
 
+if VERSION < v"1.9"
+    using Compat: stack
+end
+
 const DEBUG = Bool(parse(Int, get(ENV, "DEBUG_BIJECTORS", "0")))
 _debug(str) = @debug str