Add LKJ bijector

src/bijectors/corr.jl

@@ -10,9 +10,11 @@ scratch.
 struct CorrBijector <: Bijector{2} end
 
 function (b::CorrBijector)(x::AbstractMatrix{<:Real})    
-    w = cholesky(x).U + zero(x) # convert to dense matrix
+    w = cholesky(x).U #+ zero(x) # convert to dense matrix
     r = _link_chol_lkj(w) 
-    return r
+    return r + zero(x) # keep LowerTriangular until here can avoid some computation
+    # This dense format itself is required by a test, though I can't get the point.
+    # https://github.com/TuringLang/Bijectors.jl/blob/b0aaa98f90958a167a0b86c8e8eca9b95502c42d/test/transform.jl#L67
 end
 
 (b::CorrBijector)(X::AbstractArray{<:AbstractMatrix{<:Real}}) = map(b, X)
@@ -111,18 +113,22 @@ which is the above implementation.
 function _link_chol_lkj(w)
     K = LinearAlgebra.checksquare(w)
 
-    z = zero(w)
-
+    z = similar(w) # z is also UpperTriangular except ReverseDiff. 
+    # Some zero filling can be avoided. Though diagnoal is still needed to be filled with zero.
+
+    # This block can't be integrated with loop below, because w[1,1] != 0.
+    @inbounds z[1, 1] = 0
+
     @inbounds for j=2:K
         z[1, j] = w[1, j]
-    end
-
-    @inbounds for j=3:K, i=2:j-1
-        p = w[i, j]
-        for ip in 1:(i-1)
-            p /= sqrt(1-z[ip, j]^2)
+        z[j, j] = 0
+        for i=2:j-1
+            p = w[i, j]
+            for ip in 1:(i-1)
+                p /= sqrt(1-z[ip, j]^2)
+            end
+            z[i, j] = p
         end
-        z[i,j] = p
     end
 
     y = atanh.(z)

src/compat/reversediff.jl

@@ -237,25 +237,35 @@ _inv_link_chol_lkj(y::TrackedMatrix) = track(_inv_link_chol_lkj, y)
     return w, pullback_inv_link_chol_lkj
 end
 
-_link_chol_lkj(w::TrackedMatrix) = track(_link_chol_lkj, w)
+# zero(x) is used to convert w to dense matrix, because ReverseDiff doesn't support track
+# LowerTriangular matrix (concretely, Base.IndexStyle(::A) != Base.IndexLinear())
+# See: http://www.juliadiff.org/ReverseDiff.jl/limits/
+_link_chol_lkj(w::TrackedMatrix) = track(_link_chol_lkj, w + zero(x))
 @grad function _link_chol_lkj(w_tracked)
     w = value(w_tracked)
 
     LinearAlgebra.checksquare(w)
 
     K = size(w, 1)
-    z = zero(w)
-
-    @inbounds for j=2:K
-        z[1, j] = w[1, j]
+    z = similar(w)
+
+    # This block can't be integrated with loop below, because w[1,1] != 0.
+    @inbounds for i=1:K
+        z[i, 1] = 0
     end
 
-    @inbounds for j=3:K, i=2:j-1
-        p = w[i, j]
-        for ip in 1:(i-1)
-            p *= 1 / sqrt(1-z[ip, j]^2)
+    @inbounds for j=2:K
+        z[1, j] = w[1, j]
+        for i=j:K
+            z[i, j] = 0
+        end
+        for i=2:j-1
+            p = w[i, j]
+            for ip in 1:(i-1)
+                p /= sqrt(1-z[ip, j]^2)
+            end
+            z[i, j] = p
         end
-        z[i,j] = p
     end
 
     y = atanh.(z)

src/compat/tracker.jl

@@ -503,18 +503,21 @@ _link_chol_lkj(w::TrackedMatrix) = track(_link_chol_lkj, w)
     LinearAlgebra.checksquare(w)
 
     K = size(w, 1)
-    z = zero(w)
-
+    z = similar(w)
+
+    # This block can't be integrated with loop below, because w[1,1] != 0.
+    @inbounds z[1, 1] = 0
+
     @inbounds for j=2:K
         z[1, j] = w[1, j]
-    end
-
-    @inbounds for j=3:K, i=2:j-1
-        p = w[i, j]
-        for ip in 1:(i-1)
-            p *= 1 / sqrt(1-z[ip, j]^2)
+        z[j, j] = 0
+        for i=2:j-1
+            p = w[i, j]
+            for ip in 1:(i-1)
+                p /= sqrt(1-z[ip, j]^2)
+            end
+            z[i, j] = p
         end
-        z[i,j] = p
     end
 
     y = atanh.(z)

src/compat/zygote.jl

@@ -196,7 +196,7 @@ end
 
 @adjoint function _inv_link_chol_lkj(y)
     LinearAlgebra.checksquare(y)
-    
+
     K = size(y, 1)
 
     z = tanh.(y)
@@ -250,18 +250,21 @@ end
     LinearAlgebra.checksquare(w)
 
     K = size(w, 1)
-    z = zero(w)
-
+    z = similar(w)
+
+    # This block can't be integrated with loop below, because w[1,1] != 0.
+    @inbounds z[1, 1] = 0
+
     @inbounds for j=2:K
         z[1, j] = w[1, j]
-    end
-
-    @inbounds for j=3:K, i=2:j-1
-        p = w[i, j]
-        for ip in 1:(i-1)
-            p *= 1 / sqrt(1-z[ip, j]^2)
+        z[j, j] = 0
+        for i=2:j-1
+            p = w[i, j]
+            for ip in 1:(i-1)
+                p /= sqrt(1-z[ip, j]^2)
+            end
+            z[i, j] = p
         end
-        z[i,j] = p
     end
 
     y = atanh.(z)