Added output_length and output_size

src/bijectors/corr.jl

@@ -232,6 +232,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 +328,9 @@ function logabsdetjac(::Inverse{VecCholeskyBijector}, y::AbstractVector{<:Real})
     return _logabsdetjac_inv_chol(y)
 end
 
+output_size(::VecCholeskyBijector, sz::NTuple{2}) = output_size(VecCorrBijector(), sz)
+output_size(::Inverse{<:VecCholeskyBijector}, sz::NTuple{1}) = output_size(inverse(VecCorrBijector()), sz)
+
 """
     function _link_chol_lkj(w)
 

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,18 @@
+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}
+struct TransformedDistribution{D,B,V} <: 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
 
@@ -101,8 +99,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)

test/bijectors/corr.jl

@@ -33,6 +33,15 @@ using Bijectors: VecCorrBijector, VecCholeskyBijector, CorrBijector
         test_bijector(bvec, x; test_not_identity=d != 1, changes_of_variables_test=false)
 
         test_ad(x -> sum(bvec(bvecinv(x))), yvec)
+
+        # Check that output sizes are computed correctly.
+        tdist = transformed(dist)
+        @test length(tdist) == length(yvec)
+        @test tdist isa MultivariateDistribution
+
+        dist_unconstrained = transformed(MvNormal(zeros(length(tdist)), I), inverse(bvec))
+        @test size(dist_unconstrained) == size(x)
+        @test dist_unconstrained isa MatrixDistribution
     end
 end
 
@@ -60,6 +69,17 @@ end
             # test_bijector is commented out for now, 
             # as isapprox is not defined for ::Cholesky types (the domain of LKJCholesky)
             # test_bijector(b, x; test_not_identity=d != 1, changes_of_variables_test=false)
+
+            # Check that output sizes are computed correctly.
+            tdist = transformed(dist)
+            @test length(tdist) == length(y)
+            @test tdist isa MultivariateDistribution
+
+            dist_unconstrained = transformed(
+                MvNormal(zeros(length(tdist)), I), inverse(b)
+            )
+            @test size(dist_unconstrained) == size(x)
+            @test dist_unconstrained isa Distribution{CholeskyVariate,Continuous}
         end
     end
 end