Fix for product of Dirichlet

[torfjelde]

logpdf_with_trans currently doesn't handle stuff like products of Dirichlet

Long-term we want to remove logpdf_with_trans completely, but for now this is holding up TuringLang/DynamicPPL.jl#575 so seems worth providing a hotfix.

[torfjelde]

@yebai @sunxd3 @mhauru

[mhauru]

@devmotion has a fair comment, but the code looks good to me otherwise. I'm confused by the CI failure, the new test passes for me locally.

[torfjelde]

@devmotion good shout 👍

but the code looks good to me otherwise. I'm confused by the CI failure, the new test passes for me locally.

Seems to be a julia v1.6 thing. Looking at the deps used, I notice DistributionsAD is on 0.6.45 vs. 0.6.55 (on Julia latest). Miiight be it, though uncertain why. Lookin' into it.

EDIT: ah, it's just eachslice not supporting multiple dims on Julia v1.6 😕

[torfjelde]

The way we use eachslice is only valid on Julia 1.9 and onwards 😕

[mhauru]

There's mapslices that has been taking multiple dimensions for longer than eachslice, I think switching to using that could work, since we call map on the results of eachslice. I'm happy to give this a go given @torfjelde, unless you're already in the middle of solving it.

[torfjelde]

Go for it!

[mhauru]

mapslices doesn't accept multiple arguments the way map does, so I ended up doing some reshaping instead: #323