Fix for #255 + some other DistributionsAD-stuff (#259)

* fix for #255 and introduction of columnwise

* added some tests

* version  bump

* forgot to add the inverse

Author: torfjelde

Project.toml

@@ -1,6 +1,6 @@
 name = "Bijectors"
 uuid = "76274a88-744f-5084-9051-94815aaf08c4"
-version = "0.12.4"
+version = "0.12.5"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

src/compat/distributionsad.jl

@@ -13,8 +13,8 @@ bijector(::TuringScalMvNormal) = identity
 bijector(::TuringDiagMvNormal) = identity
 bijector(::TuringDenseMvNormal) = identity
 
-bijector(d::FillVectorOfUnivariate{Continuous}) = bijector(d.v.value)
-bijector(d::FillMatrixOfUnivariate{Continuous}) = up1(bijector(d.dists.value))
+bijector(d::FillVectorOfUnivariate{Continuous}) = elementwise(bijector(d.v.value))
+bijector(d::FillMatrixOfUnivariate{Continuous}) = elementwise(bijector(d.dists.value))
 bijector(d::MatrixOfUnivariate{Discrete}) = identity
 bijector(d::MatrixOfUnivariate{Continuous}) = TruncatedBijector(_minmax(d.dists)...)
 bijector(d::VectorOfMultivariate{Discrete}) = identity
@@ -30,7 +30,7 @@ for T in (:VectorOfMultivariate, :FillVectorOfMultivariate)
         bijector(d::$T{Continuous, <:TuringDirichlet}) = SimplexBijector()
     end
 end
-bijector(d::FillVectorOfMultivariate{Continuous}) = bijector(d.dists.value)
+bijector(d::FillVectorOfMultivariate{Continuous}) = columnwise(bijector(d.dists.value))
 
 isdirichlet(::VectorOfMultivariate{Continuous, <:Dirichlet}) = true
 isdirichlet(::VectorOfMultivariate{Continuous, <:TuringDirichlet}) = true

src/interface.jl

@@ -17,6 +17,21 @@ elementwise(f) = Base.Fix1(broadcast, f)
 # TODO: This is makes dispatching quite a bit easier, but uncertain if this is really
 # the way to go.
 elementwise(f::ComposedFunction) = ComposedFunction(elementwise(f.outer), elementwise(f.inner))
+const Columnwise{F} = Base.Fix1{typeof(eachcolmaphcat),F}
+"""
+
+Alias for `Base.Fix1(eachcolmaphcat, f)`.
+
+Represents a function `f` which is applied to each column of an input.
+"""
+columnwise(f) = Base.Fix1(eachcolmaphcat, f)
+inverse(f::Columnwise) = columnwise(inverse(f.x))
+
+transform(f::Columnwise, x::AbstractMatrix) = f(x)
+function logabsdetjac(f::Columnwise, x::AbstractMatrix)
+    return sum(Base.Fix1(logabsdetjac, f.x), eachcol(x))
+end
+with_logabsdet_jacobian(f::Columnwise, x::AbstractMatrix) = (f(x), logabsdetjac(f, x))
 
 ######################
 # Bijector interface #

test/interface.jl

@@ -196,6 +196,22 @@ end
     end
 end
 
+@testset "DistributionsAD" begin
+    @testset "$dist" for dist in [
+        filldist(Normal(), 2),
+        filldist(Normal(), 2, 3),
+        filldist(Exponential(), 2),
+        filldist(Exponential(), 2, 3),
+        filldist(filldist(Exponential(), 2), 3),
+    ]
+        x = rand(dist)
+        b = bijector(dist)
+        y = b(x)
+        td = transformed(dist)
+        @test logpdf(dist, x) - logabsdetjac(b, x) ≈ logpdf(td, y)
+    end
+end
+
 @testset "Stacked <: Bijector" begin
     # `logabsdetjac` withOUT AD
     d = Beta()