Fix incorrect output dimensions in stacked bijectors (#2066)

* Add output_length and transform dispatches for Vec bijector

This fixes #2065

* Add tests

* Bump revision

* Rename model function in simplex bijector test

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.28.2"
+version = "0.28.3"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/variational/advi.jl

@@ -6,6 +6,9 @@ end
 
 Bijectors.inverse(f::Vec) = Vec(Bijectors.inverse(f.b), f.size)
 
+Bijectors.output_length(f::Vec, sz) = Bijectors.output_length(f.b, sz)
+Bijectors.output_length(f::Vec, n::Int) = Bijectors.output_length(f.b, n)
+
 function Bijectors.with_logabsdet_jacobian(f::Vec, x)
     return Bijectors.transform(f, x), Bijectors.logabsdetjac(f, x)
 end
@@ -15,6 +18,11 @@ function Bijectors.transform(f::Vec, x::AbstractVector)
     return vec(f.b(reshape(x, f.size)))
 end
 
+function Bijectors.transform(f::Vec{N,<:Bijectors.Inverse}, x::AbstractVector) where N
+    # Reshape into shape compatible with original (forward) bijector and then `vec` again.
+    return vec(f.b(reshape(x, Bijectors.output_length(f.b.orig, prod(f.size)))))
+end
+
 function Bijectors.transform(f::Vec, x::AbstractMatrix)
     # At the moment we do batching for higher-than-1-dim spaces by simply using
     # lists of inputs rather than `AbstractArray` with `N + 1` dimension.

test/variational/advi.jl

@@ -47,4 +47,22 @@
         xs = rand(target, 10)
         @test mean(abs2, logpdf(q, xs) - logpdf(target, xs)) ≤ 0.05
     end
+
+    # regression test for:
+    # https://github.com/TuringLang/Turing.jl/issues/2065
+    @turing_testset "simplex bijector" begin
+        @model function dirichlet()
+            x ~ Dirichlet([1.0,1.0])
+            return x
+        end
+
+        m = dirichlet()
+        b = bijector(m)
+        x0 = m()
+        z0 = b(x0)
+        @test size(z0) == (1,)
+        x0_inv = inverse(b)(z0)
+        @test size(x0_inv) == size(x0)
+        @test all(x0 .≈ x0_inv)
+    end
 end