Fix for product of Dirichlet (#322)

* Added default impl of `_logabdetjac_dist` so we can support non-batch
by default

* Added test for product of `Dirichlet`

* Bump patch version

* Update src/Bijectors.jl

Co-authored-by: David Widmann <[email protected]>

* Work around eachslice limitation on Julia <v1.9 (#323)

* Work around eachslice limitation on Julia <v1.9

* Bug fix

* Fix Tapir tests

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update test/ad/chainrules.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: David Widmann <[email protected]>
Co-authored-by: Markus Hauru <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Hong Ge <[email protected]>

Project.toml

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

src/Bijectors.jl

@@ -135,6 +135,11 @@ invlink(d::Distribution, y) = inverse(bijector(d))(y)
 
 # To still allow `logpdf_with_trans` to work with "batches" in a similar way
 # as `logpdf` can.
+
+# Default catch-all so we can work with distributions by default and batch-support can be
+# added when needed.
+_logabsdetjac_dist(d::Distribution, x) = logabsdetjac(bijector(d), x)
+
 _logabsdetjac_dist(d::UnivariateDistribution, x::Real) = logabsdetjac(bijector(d), x)
 function _logabsdetjac_dist(d::UnivariateDistribution, x::AbstractArray)
     return logabsdetjac.((bijector(d),), x)

src/bijectors/product_bijector.jl

@@ -8,11 +8,10 @@ inverse(b::ProductBijector) = ProductBijector(map(inverse, b.bs))
 
 function _product_bijector_check_dim(::Val{N}, ::Val{M}) where {N,M}
     if N > M
-        throw(
-            DimensionMismatch(
-                "Number of bijectors needs to be smaller than or equal to the number of dimensions",
-            ),
-        )
+        msg = """
+        Number of bijectors needs to be smaller than or equal to the number of dimensions
+        """
+        throw(DimensionMismatch(msg))
     end
 end
 
@@ -23,7 +22,18 @@ function _product_bijector_slices(
 
     # If N < M, then the bijectors expect an input vector of dimension `M - N`.
     # To achieve this, we need to slice along the last `N` dimensions.
-    return eachslice(x; dims=ntuple(i -> i + (M - N), N))
+    slice_indices = ntuple(i -> i + (M - N), N)
+    if VERSION >= v"1.9"
+        return eachslice(x; dims=slice_indices)
+    else
+        # Earlier Julia versions can't eachslice over multiple dimensions, so reshape the
+        # slice dimensions into a single one.
+        other_dims = tuple((size(x, i) for i in 1:(M - N))...)
+        slice_dims = tuple((size(x, i) for i in (1 + M - N):M)...)
+        x_reshaped = reshape(x, other_dims..., prod(slice_dims))
+        slices = eachslice(x_reshaped; dims=M - N + 1)
+        return reshape(collect(slices), slice_dims)
+    end
 end
 
 # Specialization for case where we're just applying elementwise.

test/ad/chainrules.jl

@@ -29,14 +29,35 @@ end
 
     if @isdefined Tapir
         rng = Xoshiro(123456)
-        Tapir.TestUtils.test_rrule!!(
-            rng, Bijectors.find_alpha, x, y, z; is_primitive=true, perf_flag=:none
+        Tapir.TestUtils.test_rule(
+            rng,
+            Bijectors.find_alpha,
+            x,
+            y,
+            z;
+            is_primitive=true,
+            perf_flag=:none,
+            interp=Tapir.TapirInterpreter(),
         )
-        Tapir.TestUtils.test_rrule!!(
-            rng, Bijectors.find_alpha, x, y, 3; is_primitive=true, perf_flag=:none
+        Tapir.TestUtils.test_rule(
+            rng,
+            Bijectors.find_alpha,
+            x,
+            y,
+            3;
+            is_primitive=true,
+            perf_flag=:none,
+            interp=Tapir.TapirInterpreter(),
         )
-        Tapir.TestUtils.test_rrule!!(
-            rng, Bijectors.find_alpha, x, y, UInt32(3); is_primitive=true, perf_flag=:none
+        Tapir.TestUtils.test_rule(
+            rng,
+            Bijectors.find_alpha,
+            x,
+            y,
+            UInt32(3);
+            is_primitive=true,
+            perf_flag=:none,
+            interp=Tapir.TapirInterpreter(),
         )
     end
 

test/interface.jl

@@ -223,6 +223,15 @@ end
     end
 end
 
+@testset "ProductDistribution" begin
+    d = product_distribution(fill(Dirichlet(ones(4)), 2, 3))
+    x = rand(d)
+    b = bijector(d)
+
+    @test logpdf_with_trans(d, x, false) == logpdf(d, x)
+    @test logpdf_with_trans(d, x, true) == logpdf(d, x) - logabsdetjac(b, x)
+end
+
 @testset "DistributionsAD" begin
     @testset "$dist" for dist in [
         filldist(Normal(), 2),