Added output_length and output_size

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"

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::Tuple{Int,Int})
+    @assert sz[1] == sz[2]
+    n = sz[1]
+    return (n * (n - 1)) ÷ 2
+end
+
+function output_size(::Inverse{VecCorrBijector}, sz::Tuple{Int})
+    n = _triu1_dim_from_length(first(sz))
+    return (n, n)
+end
+
 """
     VecCholeskyBijector <: Bijector
 
@@ -317,6 +328,11 @@ function logabsdetjac(::Inverse{VecCholeskyBijector}, y::AbstractVector{<:Real})
     return _logabsdetjac_inv_chol(y)
 end
 
+output_size(::VecCholeskyBijector, sz::Tuple{Int,Int}) = output_size(VecCorrBijector(), sz)
+function output_size(::Inverse{<:VecCholeskyBijector}, sz::Tuple{Int})
+    return output_size(inverse(VecCorrBijector()), sz)
+end
+
 """
     function _link_chol_lkj(w)
 

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,13 @@ 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
+
 ######################
 # 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
 
@@ -101,8 +100,8 @@ end
 ##############################
 
 # size
-Base.length(td::Transformed) = length(td.dist)
-Base.size(td::Transformed) = size(td.dist)
+Base.length(td::Transformed) = only(output_size(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,15 @@ 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