Implementation of VecCorrBijector

src/bijectors/corr.jl

@@ -66,7 +66,7 @@ struct CorrBijector <: Bijector end
 with_logabsdet_jacobian(b::CorrBijector, x) = transform(b, x), logabsdetjac(b, x)
 
 function transform(b::CorrBijector, x::AbstractMatrix{<:Real})
-    w = cholesky(x).U  # keep LowerTriangular until here can avoid some computation
+    w = upper_triangular(parent(cholesky(x).U))  # keep LowerTriangular until here can avoid some computation
     r = _link_chol_lkj(w) 
     return r + zero(x) 
     # This dense format itself is required by a test, though I can't get the point.
@@ -75,51 +75,158 @@ end
 
 function transform(ib::Inverse{CorrBijector}, y::AbstractMatrix{<:Real})
     w = _inv_link_chol_lkj(y)
-    return w' * w
+    return pd_from_upper(w)
 end
 
-function logabsdetjac(::Inverse{CorrBijector}, y::AbstractMatrix{<:Real})
-    K = LinearAlgebra.checksquare(y)
-
-    result = float(zero(eltype(y)))
-    for j in 2:K, i in 1:(j - 1)
-        @inbounds abs_y_i_j = abs(y[i, j])
-        result += (K - i + 1) * (
-            IrrationalConstants.logtwo - (abs_y_i_j + LogExpFunctions.log1pexp(-2 * abs_y_i_j))
-        )
-    end
-
-    return result
-end
+logabsdetjac(::Inverse{CorrBijector}, Y::AbstractMatrix{<:Real}) = _logabsdetjac_chol_lkj(Y)
 function logabsdetjac(b::CorrBijector, X::AbstractMatrix{<:Real})
     #=
     It may be more efficient if we can use un-contraint value to prevent call of b
     It's recommended to directly call 
     `logabsdetjac(::Inverse{CorrBijector}, y::AbstractMatrix{<:Real})`
     if possible.
     =#
-    return -logabsdetjac(inverse(b), (b(X))) 
+    return -logabsdetjac(inverse(b), (b(X)))
 end
 
-function _inv_link_chol_lkj(y)
-    K = LinearAlgebra.checksquare(y)
+"""
+    triu_mask(X::AbstractMatrix, k::Int)
 
-    w = similar(y)
+Return a mask for elements of `X` above the `k`th diagonal.
+"""
+function triu_mask(X::AbstractMatrix, k::Int)
+    # Ensure that we're working with a square matrix.
+    LinearAlgebra.checksquare(X)
 
-    @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]
-            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
+    # Using `similar` allows us to respect device of array, etc., e.g. `CuArray`.
+    m = similar(X, Bool)
+    return triu(.~m .| m, k)
+end
+
+triu_to_vec(X::AbstractMatrix{<:Real}, k::Int) = X[triu_mask(X, k)]
+
+function update_triu_from_vec!(
+    vals::AbstractVector{<:Real},
+    k::Int,
+    X::AbstractMatrix{<:Real}
+)
+    # Ensure that we're working with one-based indexing.
+    # `triu` requires this too.
+    LinearAlgebra.require_one_based_indexing(X)
+
+    # Set the values.
+    idx = 1
+    m, n = size(X)
+    for j = 1:n
+        for i = 1:min(j - k, m)
+            X[i, j] = vals[idx]
+            idx += 1
         end
     end
-
-    return w
+
+    return X
+end
+
+function update_triu_from_vec(vals::AbstractVector{<:Real}, k::Int, dim::Int)
+    X = similar(vals, dim, dim)
+    # TODO: Do we need this?
+    X .= 0
+    return update_triu_from_vec!(vals, k, X)
+end
+
+function ChainRulesCore.rrule(::typeof(update_triu_from_vec), x::AbstractVector{<:Real}, k::Int, dim::Int)
+    function update_triu_from_vec_pullback(ΔX)
+        return (
+            ChainRulesCore.NoTangent(),
+            triu_to_vec(ChainRulesCore.unthunk(ΔX), k),
+            ChainRulesCore.NoTangent(),
+            ChainRulesCore.NoTangent()
+        )
+    end
+    return update_triu_from_vec(x, k, dim), update_triu_from_vec_pullback
+end
+
+#      n * (n - 1) / 2 = d
+# ⟺       n^2 - n - 2d = 0
+# ⟹                  n = (1 + sqrt(1 + 8d)) / 2
+_triu1_dim_from_length(d) = (1 + isqrt(1 + 8d)) ÷ 2
+
+"""
+    triu1_to_vec(X::AbstractMatrix{<:Real})
+
+Extracts elements from upper triangle of `X` with offset `1` and returns them as a vector.
+"""
+triu1_to_vec(X::AbstractMatrix) = triu_to_vec(X, 1)
+
+inverse(::typeof(triu1_to_vec)) = vec_to_triu1
+
+"""
+    vec_to_triu1(x::AbstractVector{<:Real})
+
+Constructs a matrix from a vector `x` by filling the upper triangle with offset `1`.
+"""
+function vec_to_triu1(x::AbstractVector)
+    n = _triu1_dim_from_length(length(x))
+    X = update_triu_from_vec(x, 1, n)
+    return upper_triangular(X)
+end
+
+inverse(::typeof(vec_to_triu1)) = triu1_to_vec
+
+"""
+    VecCorrBijector <: Bijector
+
+Similar to `CorrBijector`, but correlation matrix to a vector,
+and its inverse transforms vector to a correlation matrix.
+
+See also: [`CorrBijector`](@ref)
+
+# Example
+
+```jldoctest
+julia> using LinearAlgebra
+
+julia> using StableRNGs; rng = StableRNG(42);
+
+julia> b = Bijectors.VecCorrBijector();
+
+julia> X = rand(rng, LKJ(3, 1))  # Sample a correlation matrix.
+3×3 Matrix{Float64}:
+  1.0       -0.705273   -0.348638
+ -0.705273   1.0         0.0534538
+ -0.348638   0.0534538   1.0
+
+julia> y = b(X)  # Transform to unconstrained vector representation.
+3-element Vector{Float64}:
+ -0.8777149781928181
+ -0.3638927608636788
+ -0.29813769428942216
+
+julia> inverse(b)(y) ≈ X  # (✓) Round-trip through `b` and its inverse.
+true
+"""
+struct VecCorrBijector <: Bijector end
+with_logabsdet_jacobian(b::VecCorrBijector, x) = transform(b, x), logabsdetjac(b, x)
+
+function transform(::VecCorrBijector, X::AbstractMatrix{<:Real})
+    w = upper_triangular(parent(cholesky(X).U))
+    r = _link_chol_lkj(w)
+
+    # Extract only the upper triangle of `r`.
+    return triu1_to_vec(r)
+end
+
+function transform(::Inverse{VecCorrBijector}, y::AbstractVector{<:Real})
+    Y = vec_to_triu1(y)
+    w = _inv_link_chol_lkj(Y)
+    return pd_from_upper(w)
+end
+
+function logabsdetjac(b::VecCorrBijector, x)
+    return -logabsdetjac(inverse(b), b(x))
+end
+function logabsdetjac(::Inverse{VecCorrBijector}, y::AbstractVector{<:Real})
+    return _logabsdetjac_chol_lkj(vec_to_triu1(y))
 end
 
 """
@@ -138,7 +245,7 @@ end
 But this implementation will not work when w[i-1, j] = 0.
 Though it is a zero measure set, unit matrix initialization will not work.
 
-For equivelence, following explanations is given by @torfjelde:
+For equivalence, following explanations is given by @torfjelde:
 
 For `(i, j)` in the loop below, we define
 
@@ -155,6 +262,7 @@ and so
 which is the above implementation.
 """
 function _link_chol_lkj(w)
+    # TODO: Implement adjoint to support reverse-mode AD backends properly.
     K = LinearAlgebra.checksquare(w)
 
     z = similar(w) # z is also UpperTriangular. 
@@ -163,16 +271,56 @@ function _link_chol_lkj(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
+    @inbounds for j = 2:K
         z[1, j] = atanh(w[1, j])
         tmp = sqrt(1 - w[1, j]^2)
-        for i in 2:(j - 1)
+        for i in 2:(j-1)
             p = w[i, j] / tmp
             tmp *= sqrt(1 - p^2)
             z[i, j] = atanh(p)
         end
         z[j, j] = 0
     end
-    
+
     return z
 end
+
+"""
+    _inv_link_chol_lkj(y)
+
+Inverse link function for cholesky factor.
+"""
+function _inv_link_chol_lkj(y)
+    # TODO: Implement adjoint to support reverse-mode AD backends properly.
+    K = LinearAlgebra.checksquare(y)
+
+    w = 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]
+            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
+
+    return w
+end
+
+function _logabsdetjac_chol_lkj(Y::AbstractMatrix)
+    K = LinearAlgebra.checksquare(Y)
+
+    result = float(zero(eltype(Y)))
+    for j in 2:K, i in 1:(j-1)
+        @inbounds abs_y_i_j = abs(Y[i, j])
+        result += (K - i + 1) * (
+            IrrationalConstants.logtwo - (abs_y_i_j + LogExpFunctions.log1pexp(-2 * abs_y_i_j))
+        )
+    end
+    return result
+end

src/bijectors/pd.jl

@@ -9,16 +9,14 @@ function replace_diag(f, X)
 end
 transform(b::PDBijector, X::AbstractMatrix{<:Real}) = pd_link(X)
 function pd_link(X)
-    Y = lower(parent(cholesky(X; check = true).L))
+    Y = lower_triangular(parent(cholesky(X; check = true).L))
     return replace_diag(log, Y)
 end
-lower(A::AbstractMatrix) = convert(typeof(A), LowerTriangular(A))
 
 function transform(ib::Inverse{PDBijector}, Y::AbstractMatrix{<:Real})
     X = replace_diag(exp, Y)
-    return getpd(X)
+    return pd_from_lower(X)
 end
-getpd(X) = LowerTriangular(X) * LowerTriangular(X)'
 
 function logabsdetjac(b::PDBijector, X::AbstractMatrix{<:Real})
     T = eltype(X)

src/compat/reversediff.jl

@@ -1,13 +1,14 @@
 module ReverseDiffCompat
 
 using ..ReverseDiff: ReverseDiff, @grad, value, track, TrackedReal, TrackedVector, 
-    TrackedMatrix
+    TrackedMatrix, @grad_from_chainrules
 using Requires, LinearAlgebra
 
 using ..Bijectors: Elementwise, SimplexBijector, maphcat, simplex_link_jacobian, 
     simplex_invlink_jacobian, simplex_logabsdetjac_gradient, Inverse
 import ..Bijectors: _eps, logabsdetjac, _logabsdetjac_scale, _simplex_bijector, 
-    _simplex_inv_bijector, replace_diag, jacobian, getpd, lower, 
+    _simplex_inv_bijector, replace_diag, jacobian, pd_from_lower, pd_from_upper,
+    lower_triangular, upper_triangular,
     _inv_link_chol_lkj, _link_chol_lkj, _transform_ordered, _transform_inverse_ordered,
     find_alpha
 
@@ -136,18 +137,34 @@ logabsdetjac(b::SimplexBijector, x::Union{TrackedVector, TrackedMatrix}) = track
     end
 end
 
-getpd(X::TrackedMatrix) = track(getpd, X)
-@grad function getpd(X::AbstractMatrix)
+pd_from_lower(X::TrackedMatrix) = track(pd_from_lower, X)
+@grad function pd_from_lower(X::AbstractMatrix)
     Xd = value(X)
     return LowerTriangular(Xd) * LowerTriangular(Xd)', Δ -> begin
         Xl = LowerTriangular(Xd)
         return (LowerTriangular(Δ' * Xl + Δ * Xl),)
     end
 end
-lower(A::TrackedMatrix) = track(lower, A)
-@grad function lower(A::AbstractMatrix)
+
+pd_from_upper(X::TrackedMatrix) = track(pd_from_upper, X)
+@grad function pd_from_upper(X::AbstractMatrix)
+    Xd = value(X)
+    return UpperTriangular(Xd)' * UpperTriangular(Xd), Δ -> begin
+        Xu = UpperTriangular(Xd)
+        return (UpperTriangular(Δ * Xu + Δ' * Xu),)
+    end
+end
+
+lower_triangular(A::TrackedMatrix) = track(lower_triangular, A)
+@grad function lower_triangular(A::AbstractMatrix)
     Ad = value(A)
-    return lower(Ad), Δ -> (lower(Δ),)
+    return lower_triangular(Ad), Δ -> (lower_triangular(Δ),)
+end
+
+upper_triangular(A::TrackedMatrix) = track(upper_triangular, A)
+@grad function upper_triangular(A::AbstractMatrix)
+    Ad = value(A)
+    return upper_triangular(Ad), Δ -> (upper_triangular(Δ),)
 end
 
 function find_alpha(wt_y::T, wt_u_hat::T, b::T) where {T<:TrackedReal}
@@ -181,6 +198,8 @@ end
     return y, (wrap_chainrules_output ∘ Base.tail ∘ dy)
 end
 
+@grad_from_chainrules update_triu_from_vec(vals::TrackedVector{<:Real}, k::Int, dim::Int)
+
 # NOTE: Probably doesn't work in complete generality.
 wrap_chainrules_output(x) = x
 wrap_chainrules_output(x::ChainRulesCore.AbstractZero) = nothing

src/compat/tracker.jl

@@ -281,19 +281,19 @@ end
 (b::Elementwise{typeof(log)})(x::TrackedVector) = log.(x)::vectorof(float(eltype(x)))
 (b::Elementwise{typeof(log)})(x::TrackedMatrix) = log.(x)::matrixof(float(eltype(x)))
 
-Bijectors.getpd(X::TrackedMatrix) = track(Bijectors.getpd, X)
-@grad function Bijectors.getpd(X::AbstractMatrix)
+Bijectors.pd_from_lower(X::TrackedMatrix) = track(Bijectors.pd_from_lower, X)
+@grad function Bijectors.pd_from_lower(X::AbstractMatrix)
     Xd = data(X)
     return Bijectors.LowerTriangular(Xd) * Bijectors.LowerTriangular(Xd)', Δ -> begin
         Xl = Bijectors.LowerTriangular(Xd)
         return (Bijectors.LowerTriangular(Δ' * Xl + Δ * Xl),)
     end
 end
 
-Bijectors.lower(A::TrackedMatrix) = track(Bijectors.lower, A)
-@grad function Bijectors.lower(A::AbstractMatrix)
+Bijectors.lower_triangular(A::TrackedMatrix) = track(Bijectors.lower_triangular, A)
+@grad function Bijectors.lower_triangular(A::AbstractMatrix)
     Ad = data(A)
-    return Bijectors.lower(Ad), Δ -> (Bijectors.lower(Δ),)
+    return Bijectors.lower_triangular(Ad), Δ -> (Bijectors.lower_triangular(Δ),)
 end
 
 Bijectors._inv_link_chol_lkj(y::TrackedMatrix) = track(Bijectors._inv_link_chol_lkj, y)

src/compat/zygote.jl

@@ -158,10 +158,10 @@ end
     end
     return pullback(_maximum, d)
 end
-@adjoint function lower(A::AbstractMatrix)
-    return lower(A), Δ -> (lower(Δ),)
+@adjoint function lower_triangular(A::AbstractMatrix)
+    return lower_triangular(A), Δ -> (lower_triangular(Δ),)
 end
-@adjoint function getpd(X::AbstractMatrix)
+@adjoint function pd_from_lower(X::AbstractMatrix)
     return LowerTriangular(X) * LowerTriangular(X)', Δ -> begin
         Xl = LowerTriangular(X)
         return (LowerTriangular(Δ' * Xl + Δ * Xl),)