Vector-version for PDBijector

Project.toml

@@ -1,6 +1,6 @@
 name = "Bijectors"
 uuid = "76274a88-744f-5084-9051-94815aaf08c4"
-version = "0.12.8"
+version = "0.13.0"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

ext/BijectorsTrackerExt.jl

@@ -529,7 +529,26 @@ end
 
 # NOTE: Probably doesn't work in complete generality.
 wrap_chainrules_output(x) = x
+wrap_chainrules_output(x::ChainRulesCore.Thunk) = wrap_chainrules_output(ChainRulesCore.unthunk(x))
 wrap_chainrules_output(x::ChainRulesCore.AbstractZero) = nothing
 wrap_chainrules_output(x::Tuple) = map(wrap_chainrules_output, x)
 
+# `update_triu_from_vec`
+function Bijectors.update_triu_from_vec(vals::TrackedVector{<:Real}, k::Int, dim::Int)
+    return track(Bijectors.update_triu_from_vec, vals, k, dim)
+end
+
+@grad function Bijectors.update_triu_from_vec(vals::TrackedVector{<:Real}, k::Int, dim::Int)
+    # HACK: This doesn't support higher order!
+    y, dy = ChainRulesCore.rrule(Bijectors.update_triu_from_vec, data(vals), k, dim)
+    return y, (wrap_chainrules_output ∘ Base.tail ∘ dy)
+end
+
+Bijectors.upper_triangular(A::TrackedMatrix) = track(Bijectors.upper_triangular, A)
+@grad function Bijectors.upper_triangular(A::AbstractMatrix)
+    Ad = data(A)
+    return Bijectors.upper_triangular(Ad), Δ -> (Bijectors.upper_triangular(Δ),)
+end
+
+
 end

src/bijectors/corr.jl

@@ -89,97 +89,6 @@ function logabsdetjac(b::CorrBijector, X::AbstractMatrix{<:Real})
     return -logabsdetjac(inverse(b), (b(X)))
 end
 
-"""
-    triu_mask(X::AbstractMatrix, k::Int)
-
-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)
-
-    # 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 in 1:n
-        for i in 1:min(j - k, m)
-            X[i, j] = vals[idx]
-            idx += 1
-        end
-    end
-
-    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
-
-function vec_to_triu1_row_index(idx)
-    # Assumes that vector was saved in a column-major order
-    # and that vector is one-based indexed.
-    M = _triu1_dim_from_length(idx - 1)
-    return idx - (M * (M - 1) ÷ 2)
-end
-
 """
     VecCorrBijector <: Bijector
 
@@ -232,6 +141,17 @@ function logabsdetjac(::Inverse{VecCorrBijector}, y::AbstractVector{<:Real})
     return _logabsdetjac_inv_corr(y)
 end
 
+function output_size(::VecCorrBijector, sz::NTuple{2})
+    @assert sz[1] == sz[2]
+    n = sz[1]
+    return (n * (n - 1)) ÷ 2
+end
+
+function output_size(::Inverse{VecCorrBijector}, sz::NTuple{1})
+    n = _triu1_dim_from_length(first(sz))
+    return (n, n)
+end
+
 """
     VecCholeskyBijector <: Bijector
 
@@ -317,6 +237,11 @@ function logabsdetjac(::Inverse{VecCholeskyBijector}, y::AbstractVector{<:Real})
     return _logabsdetjac_inv_chol(y)
 end
 
+output_size(::VecCholeskyBijector, sz::NTuple{2}) = output_size(VecCorrBijector(), sz)
+function output_size(::Inverse{<:VecCholeskyBijector}, sz::NTuple{1})
+    return output_size(inverse(VecCorrBijector()), sz)
+end
+
 """
     function _link_chol_lkj(w)
 

src/bijectors/pd.jl

@@ -40,3 +40,48 @@ end
 function with_logabsdet_jacobian(b::PDBijector, X)
     return transform(b, X), logabsdetjac(b, X)
 end
+
+struct PDVecBijector <: Bijector end
+
+function _triu_dim_from_length(d)
+    # (n^2 + n) / 2 = d
+    #   n² + n - 2d = 0
+    #             n = (-1 + sqrt(1 + 8d)) / 2
+    return (-1 + isqrt(1 + 8 * d)) ÷ 2
+end
+
+"""
+    vec_to_triu(x::AbstractVector{<:Real})
+
+Constructs a matrix from a vector `x` by filling the upper triangle.
+"""
+function vec_to_triu(x::AbstractVector)
+    n = _triu_dim_from_length(length(x))
+    X = update_triu_from_vec(x, 0, n)
+    return upper_triangular(X)
+end
+
+transform(::PDVecBijector, X::AbstractMatrix{<:Real}) = pd_vec_link(X)
+pd_vec_link(X) = triu_to_vec(transpose(pd_link(X)), 0)
+
+function transform(::Inverse{PDVecBijector}, y::AbstractVector{<:Real})
+    Y = permutedims(vec_to_triu(y))
+    return transform(inverse(PDBijector()), Y)
+end
+
+logabsdetjac(::PDVecBijector, X::AbstractMatrix{<:Real}) = logabsdetjac(PDBijector(), X)
+
+function with_logabsdet_jacobian(b::PDVecBijector, X)
+    return transform(b, X), logabsdetjac(b, X)
+end
+
+function output_size(::PDVecBijector, sz::Tuple{Int,Int})
+    n = first(sz)
+    d = (n^2 + n) ÷ 2
+    return (d,)
+end
+
+function output_size(::Inverse{PDVecBijector}, sz::Tuple{Int})
+    n = _triu_dim_from_length(first(sz))
+    return (n, n)
+end

src/bijectors/reshape.jl

@@ -25,3 +25,5 @@ end
 inverse(b::Reshape) = Reshape(b.out_shape, b.in_shape)
 
 with_logabsdet_jacobian(b::Reshape, x) = reshape(x, b.out_shape), zero(eltype(x))
+
+output_size(b::Reshape, in_size) = b.out_shape

src/interface.jl

@@ -35,6 +35,26 @@ function logabsdetjac(f::Columnwise, x::AbstractMatrix)
 end
 with_logabsdet_jacobian(f::Columnwise, x::AbstractMatrix) = (f(x), logabsdetjac(f, x))
 
+"""
+    output_size(f, sz)
+
+Returns the output size of `f` given the input size `sz`.
+"""
+output_size(f, sz) = sz
+
+"""
+    output_length(f, len::Int)
+    output_length(f, sz::Tuple)
+
+Returns the output length of `f` given the input length `len` or size `sz`.
+"""
+output_length(f, len::Int) = len
+function output_length(f, len::Tuple)
+    sz = output_size(f, len)
+    @assert length(sz) == 1
+    return first(sz)
+end
+
 ######################
 # Bijector interface #
 ######################

src/transformed_distribution.jl

@@ -1,20 +1,19 @@
+function variateform(d::Distribution, b)
+    sz_in = size(d)
+    sz_out = output_size(b, sz_in)
+    return ArrayLikeVariate{length(sz_out)}
+end
+
+variateform(::MultivariateDistribution, ::Inverse{VecCholeskyBijector}) = CholeskyVariate
+
 # Transformed distributions
 struct TransformedDistribution{D,B,V} <:
-       Distribution{V,Continuous} where {D<:Distribution{V,Continuous},B}
+       Distribution{V,Continuous} where {D<:ContinuousDistribution,B}
     dist::D
     transform::B
 
-    function TransformedDistribution(d::UnivariateDistribution, b)
-        return new{typeof(d),typeof(b),Univariate}(d, b)
-    end
-    function TransformedDistribution(d::MultivariateDistribution, b)
-        return new{typeof(d),typeof(b),Multivariate}(d, b)
-    end
-    function TransformedDistribution(d::MatrixDistribution, b)
-        return new{typeof(d),typeof(b),Matrixvariate}(d, b)
-    end
-    function TransformedDistribution(d::Distribution{CholeskyVariate}, b)
-        return new{typeof(d),typeof(b),CholeskyVariate}(d, b)
+    function TransformedDistribution(d::ContinuousDistribution, b)
+        return new{typeof(d),typeof(b),variateform(d, b)}(d, b)
     end
 end
 
@@ -83,8 +82,8 @@ bijector(d::BoundedDistribution) = bijector_bounded(d)
 const LowerboundedDistribution = Union{Pareto,Levy}
 bijector(d::LowerboundedDistribution) = bijector_lowerbounded(d)
 
-bijector(d::PDMatDistribution) = PDBijector()
-bijector(d::MatrixBeta) = PDBijector()
+bijector(d::PDMatDistribution) = PDVecBijector()
+bijector(d::MatrixBeta) = PDVecBijector()
 
 bijector(d::LKJ) = VecCorrBijector()
 bijector(d::LKJCholesky) = VecCholeskyBijector(d.uplo)
@@ -101,8 +100,8 @@ end
 ##############################
 
 # size
-Base.length(td::Transformed) = length(td.dist)
-Base.size(td::Transformed) = size(td.dist)
+Base.length(td::Transformed) = output_length(td.transform, size(td.dist))
+Base.size(td::Transformed) = output_size(td.transform, size(td.dist))
 
 function logpdf(td::UnivariateTransformed, y::Real)
     x, logjac = with_logabsdet_jacobian(inverse(td.transform), y)

src/utils.jl

@@ -18,3 +18,94 @@ cholesky_factor(X::AbstractMatrix) = cholesky_factor(cholesky(Hermitian(X)))
 cholesky_factor(X::Cholesky) = X.U
 cholesky_factor(X::UpperTriangular) = X
 cholesky_factor(X::LowerTriangular) = X
+
+"""
+    triu_mask(X::AbstractMatrix, k::Int)
+
+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)
+
+    # 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 in 1:n
+        for i in 1:min(j - k, m)
+            X[i, j] = vals[idx]
+            idx += 1
+        end
+    end
+
+    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
+
+function vec_to_triu1_row_index(idx)
+    # Assumes that vector was saved in a column-major order
+    # and that vector is one-based indexed.
+    M = _triu1_dim_from_length(idx - 1)
+    return idx - (M * (M - 1) ÷ 2)
+end