filldist on distributions requiring SimplexBijector

[torfjelde]

MWE:

julia> @model demo() = x ~ filldist(Dirichlet(ones(2)), 3)
demo (generic function with 2 methods)

julia> sample(demo(), NUTS(), 1000)
Sampling 100%|███████████████████████████████████████████████| Time: 0:00:02
ERROR: StackOverflowError:
Stacktrace:
     [1] Array
       @ ./boot.jl:479 [inlined]
     [2] Array
       @ ./boot.jl:487 [inlined]
     [3] similar
       @ ./array.jl:421 [inlined]
     [4] _simplex_bijector
       @ ~/.julia/packages/Bijectors/gR4QV/src/bijectors/simplex.jl:22 [inlined]
     [5] transform
       @ ~/.julia/packages/Bijectors/gR4QV/src/bijectors/simplex.jl:16 [inlined]
     [6] with_logabsdet_jacobian(b::Bijectors.SimplexBijector, x::Matrix{Float64})
       @ Bijectors ~/.julia/packages/Bijectors/gR4QV/src/bijectors/simplex.jl:14
     [7] logabsdetjac(b::Bijectors.SimplexBijector, x::Matrix{Float64})
       @ Bijectors ~/.julia/packages/Bijectors/gR4QV/src/interface.jl:119
--- the last 2 lines are repeated 39989 more times ---
 [79986] with_logabsdet_jacobian(b::Bijectors.SimplexBijector, x::Matrix{Float64})
       @ Bijectors ~/.julia/packages/Bijectors/gR4QV/src/bijectors/simplex.jl:14

Fix:

# HACK: Avoids method ambiguity in `logabsdetjac`.
function Bijectors.logabsdetjac(b::Bijectors.SimplexBijector, x::AbstractMatrix{<:Real})
    return sum(Base.Fix1(logabsdetjac, b), eachcol(x))
end

# HACK: Fixes incorrect shape handling for the inverse map of `SimplexBijector`.
function DynamicPPL.with_logabsdet_jacobian_and_reconstruct(
    f::Bijectors.Inverse{<:Bijectors.SimplexBijector},
    dist,
    y
)
    (d, ns...) = size(dist)
    yreshaped = reshape(y, d - 1, ns...)
    x, logjac = with_logabsdet_jacobian(f, yreshaped)
    return x, logjac
end

Will make PRs later today / tomorrow!

[torfjelde]

Should now be fixed once a new release of DPPL has been made 👍