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 <[email protected]> * fixed typo --------- Co-authored-by: David Widmann <[email protected]> Co-authored-by: Hong Ge <[email protected]>
TuringLang · May 25, 2023 · 529f6f6 · 529f6f6 · torfjelde · May 25, 2023
1 parent 2df5699
commit 529f6f6
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 3 deletions.
diff --git a/Project.toml b/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"

diff --git a/src/bijectors/reshape.jl b/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))
diff --git a/src/interface.jl b/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")

diff --git a/src/transformed_distribution.jl b/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 #
 ##############################

diff --git a/test/bijectors/reshape.jl b/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
diff --git a/test/interface.jl b/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

diff --git a/test/runtests.jl b/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"