Support for Distributions.ProductNamedTupleDistribution

[vboussange]

Hey there,
Would it be conceivable to have something like

using Turing
import Distributions: ProductNamedTupleDistribution, Uniform, InverseGamma

params_priors = ProductNamedTupleDistribution((m = Uniform(0, 1), s² = InverseGamma(2, 3)))

@model function gdemo(x)
    params ~ params_priors
    x ~ Normal(params.m, sqrt(params.s²))
end

chain = sample(gdemo(1.2), NUTS(), 1000, progress=false)

supported in Turing.jl at some point? This could be very helpful to handle e.g. Lux.jl models.

[penelopeysm]

Hello @vboussange! Thanks for opening this issue. Yup that's on our to-do list: TuringLang/Bijectors.jl#388 so it will definitely be done at some point. Can't promise when, but this issue will help us prioritise it :)

[yebai]

This is also a useful feature for embedding JuliaBUGS models inside Turing. See TuringLang/JuliaBUGS.jl#27

Additionally, SSMProblems can be compiled into Distributions.ProductNamedTupleDistribution too, I think.

[vboussange]

Hey @penelopeysm, great that you have this in the pipeline. As a side note, I think that ProductNamedTupleDistribution does not support nested tuples, but this would helpful, at least in the context of Lux models.

[penelopeysm]

As a side note, I think that ProductNamedTupleDistribution does not support nested tuples, but this would helpful, at least in the context of Lux models.

Hmm, I can get it to do nested NamedTuples like this. Did you have something else in mind?

julia> using Distributions

julia> n = Normal()
Normal{Float64}(μ=0.0, σ=1.0)

julia> d = product_distribution((x=n, y=n))
ProductNamedTupleDistribution{(:x, :y)}(
x: Normal{Float64}(μ=0.0, σ=1.0)
y: Normal{Float64}(μ=0.0, σ=1.0)
)


julia> d2 = product_distribution((a=d, y=d))
ProductNamedTupleDistribution{(:a, :y)}(
a: ProductNamedTupleDistribution{(:x, :y)}(
x: Normal{Float64}(μ=0.0, σ=1.0)
y: Normal{Float64}(μ=0.0, σ=1.0)
)

y: ProductNamedTupleDistribution{(:x, :y)}(
x: Normal{Float64}(μ=0.0, σ=1.0)
y: Normal{Float64}(μ=0.0, σ=1.0)
)

)


julia> rand(d2)
(a = (x = 0.8357747215330494, y = 0.9587513455948373), y = (x = 0.8237375628453165, y = 1.285369184877658))

[vboussange]

Oh right, no, this looks good!

[penelopeysm]

TuringLang/Bijectors.jl#418 will fix this (when it's merged / released)

[penelopeysm]

Spoke too soon. Well, that's the first step. There will be more machinery needed in DynamicPPL.

[penelopeysm]

v0.41 should support ProductNamedTupleDistribution! If you run into any problems please feel free to ping / open a new issue :)