Added Reshape (#257)

* initial work

* added tests for reshape

* bump patch version

* added tests for some reshaped distribtuions

* minor change

* added docstring to Reshape

* bump distributions version

* implement with_logabsdet_jacobian instead

* Update src/bijectors/reshape.jl

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>

* fixed typo

---------

Co-authored-by: David Widmann <devmotion@users.noreply.github.com>
Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

Project.toml

@@ -26,7 +26,7 @@ ArgCheck = "1, 2"
 ChainRulesCore = "0.10.11, 1"
 ChangesOfVariables = "0.1"
 Compat = "3, 4"
-Distributions = "0.23.3, 0.24, 0.25"
+Distributions = "0.25.33"
 Functors = "0.1, 0.2, 0.3, 0.4"
 InverseFunctions = "0.1"
 IrrationalConstants = "0.1, 0.2"

src/bijectors/reshape.jl

@@ -0,0 +1,27 @@
+"""
+    Reshape(in_shape, out_shape)
+
+A [`Bijector`](@ref) that reshapes the input to the output shape.
+
+# Example
+
+```jldoctest
+julia> using Bijectors: Reshape
+
+julia> b = Reshape((2, 3), (3, 2))
+Reshape{Tuple{Int64, Int64}, Tuple{Int64, Int64}}((2, 3), (3, 2))
+
+julia> Array(transform(b, reshape(1:6, 2, 3)))
+3×2 Matrix{Int64}:
+ 1  4
+ 2  5
+ 3  6
+"""
+struct Reshape{N1,N2} <: Bijector
+    in_shape::NTuple{N1,Int}
+    out_shape::NTuple{N2,Int}
+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))

src/interface.jl

@@ -200,6 +200,7 @@ logabsdetjac!(::typeof(identity), x, logjac) = logjac
 # General
 include("bijectors/composed.jl")
 include("bijectors/stacked.jl")
+include("bijectors/reshape.jl")
 
 # Specific
 include("bijectors/exp_log.jl")

src/transformed_distribution.jl

@@ -79,6 +79,13 @@ bijector(d::MatrixBeta) = PDBijector()
 
 bijector(d::LKJ) = CorrBijector()
 
+function bijector(d::Distributions.ReshapedDistribution)
+    inner_dims = size(d.dist)
+    outer_dims = d.dims
+    b = Reshape(outer_dims, inner_dims)
+    return inverse(b) ∘ bijector(d.dist) ∘ b
+end
+
 ##############################
 # Distributions.jl interface #
 ##############################

test/bijectors/reshape.jl

@@ -0,0 +1,9 @@
+using Bijectors: Reshape
+
+@testset "Reshape" begin
+    dist = reshape(product_distribution(fill(InverseGamma(2, 3), 10)), 2, 5)
+    b = bijector(dist)
+
+    x = rand(dist)
+    test_bijector(b, x)
+end

test/interface.jl

@@ -120,7 +120,8 @@ end
         Dirichlet([1000 * one(Float64), eps(Float64)]), 
         Dirichlet([eps(Float64), 1000 * one(Float64)]),
         transformed(MvNormal(randn(10), Diagonal(exp.(randn(10))))),
-        transformed(MvLogNormal(MvNormal(randn(10), Diagonal(exp.(randn(10))))))
+        transformed(MvLogNormal(MvNormal(randn(10), Diagonal(exp.(randn(10)))))),
+        transformed(reshape(product_distribution(fill(InverseGamma(2, 3), 6)), 2, 3)),
     ]
 
     for dist in vector_dists
@@ -172,7 +173,8 @@ end
         InverseWishart(v,S),
         TuringWishart(v,S),
         TuringInverseWishart(v,S),
-        LKJ(3, 1.)
+        LKJ(3, 1.),
+        reshape(MvNormal(zeros(6), I), 2, 3),
     ]
 
     for dist in matrix_dists

test/runtests.jl

@@ -36,6 +36,7 @@ if GROUP == "All" || GROUP == "Interface"
     include("bijectors/coupling.jl")
     include("bijectors/ordered.jl")
     include("bijectors/pd.jl")
+    include("bijectors/reshape.jl")
 end
 
 if GROUP == "All" || GROUP == "AD"