diff --git a/Project.toml b/Project.toml index 953cd6733..04502a0bc 100644 --- a/Project.toml +++ b/Project.toml @@ -12,6 +12,7 @@ BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" +DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" @@ -51,6 +52,7 @@ Bijectors = "0.13.9" ChainRulesCore = "1" Compat = "4" ConstructionBase = "1.5.4" +DensityInterface = "0.4" Distributions = "0.25" DocStringExtensions = "0.9" EnzymeCore = "0.6, 0.7" diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl index 0d934a09a..764d1ba3d 100644 --- a/src/DynamicPPL.jl +++ b/src/DynamicPPL.jl @@ -10,6 +10,7 @@ using OrderedCollections: OrderedDict using AbstractMCMC: AbstractMCMC using ADTypes: ADTypes using BangBang: BangBang, push!!, empty!!, setindex!! +using DensityInterface: DensityInterface using MacroTools: MacroTools using ConstructionBase: ConstructionBase using Accessors: Accessors diff --git a/src/compiler.jl b/src/compiler.jl index 898acad10..5839e53bb 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -169,15 +169,18 @@ isliteral(e::Expr) = !isempty(e.args) && all(isliteral, e.args) Check if the right-hand side `x` of a `~` is a `Distribution` or an array of `Distributions`, then return `x`. """ -function check_tilde_rhs(@nospecialize(x)) +check_tilde_rhs(x) = check_tilde_rhs(x, DensityInterface.DensityKind(x)) +check_tilde_rhs(x::Distribution) = x +check_tilde_rhs(x::AbstractArray{<:Distribution}) = x +# check_tilde_rhs(x::AbstractArray{T}) where {T} = check_tilde_rhs(x, DensityInterface.DensityKind(T)) +check_tilde_rhs(x, ::DensityInterface.IsOrHasDensity) = x +function check_tilde_rhs(@nospecialize(x), ::DensityInterface.NoDensity) return throw( ArgumentError( "the right-hand side of a `~` must be a `Distribution` or an array of `Distribution`s", ), ) end -check_tilde_rhs(x::Distribution) = x -check_tilde_rhs(x::AbstractArray{<:Distribution}) = x """ unwrap_right_vn(right, vn)