Merge remote-tracking branch 'origin/master' into release-0.35

Author: mhauru

Project.toml

@@ -15,6 +15,9 @@ ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+# TODO(penelopeysm,mhauru) KernelAbstractions is only a dependency so that we can pin its version, see
+# https://github.com/TuringLang/DynamicPPL.jl/pull/781#event-16017866767
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
@@ -55,6 +58,9 @@ Compat = "4"
 ConstructionBase = "1.5.4"
 Distributions = "0.25"
 DocStringExtensions = "0.9"
+# TODO(penelopeysm,mhauru) See https://github.com/TuringLang/DynamicPPL.jl/pull/781#event-16017866767
+# for why KernelAbstractions is pinned like this.
+KernelAbstractions = "< 0.9.32"
 EnzymeCore = "0.6 - 0.8"
 ForwardDiff = "0.10"
 JET = "0.9"

src/values_as_in_model.jl

@@ -19,9 +19,9 @@ wants to extract the realization of a model in a constrained space.
 # Fields
 $(TYPEDFIELDS)
 """
-struct ValuesAsInModelContext{T,C<:AbstractContext} <: AbstractContext
+struct ValuesAsInModelContext{C<:AbstractContext} <: AbstractContext
     "values that are extracted from the model"
-    values::T
+    values::OrderedDict
     "whether to extract variables on the LHS of :="
     include_colon_eq::Bool
     "child context"
@@ -114,34 +114,32 @@ function dot_tilde_assume(
 end
 
 """
-    values_as_in_model(model::Model, include_colon_eq::Bool[, varinfo::AbstractVarInfo, context::AbstractContext])
-    values_as_in_model(rng::Random.AbstractRNG, model::Model, include_colon_eq::Bool[, varinfo::AbstractVarInfo, context::AbstractContext])
+    values_as_in_model(model::Model, include_colon_eq::Bool, varinfo::AbstractVarInfo[, context::AbstractContext])
 
 Get the values of `varinfo` as they would be seen in the model.
 
-If no `varinfo` is provided, then this is effectively the same as
-[`Base.rand(rng::Random.AbstractRNG, model::Model)`](@ref).
+More specifically, this method attempts to extract the realization _as seen in
+the model_. For example, `x[1] ~ truncated(Normal(); lower=0)` will result in a
+realization that is compatible with `truncated(Normal(); lower=0)` -- i.e. one
+where the value of `x[1]` is positive -- regardless of whether `varinfo` is
+working in unconstrained space.
 
-More specifically, this method attempts to extract the realization _as seen in the model_.
-For example, `x[1] ~ truncated(Normal(); lower=0)` will result in a realization compatible
-with `truncated(Normal(); lower=0)` regardless of whether `varinfo` is working in unconstrained
-space.
-
-Hence this method is a "safe" way of obtaining realizations in constrained space at the cost
-of additional model evaluations.
+Hence this method is a "safe" way of obtaining realizations in constrained
+space at the cost of additional model evaluations.
 
 # Arguments
 - `model::Model`: model to extract realizations from.
 - `include_colon_eq::Bool`: whether to also include variables on the LHS of `:=`.
 - `varinfo::AbstractVarInfo`: variable information to use for the extraction.
-- `context::AbstractContext`: context to use for the extraction. If `rng` is specified, then `context`
-    will be wrapped in a [`SamplingContext`](@ref) with the provided `rng`.
+- `context::AbstractContext`: base context to use for the extraction. Defaults
+   to `DynamicPPL.DefaultContext()`.
 
 # Examples
 
 ## When `VarInfo` fails
 
-The following demonstrates a common pitfall when working with [`VarInfo`](@ref) and constrained variables.
+The following demonstrates a common pitfall when working with [`VarInfo`](@ref)
+and constrained variables.
 
 ```jldoctest
 julia> using Distributions, StableRNGs
@@ -191,19 +189,10 @@ true
 function values_as_in_model(
     model::Model,
     include_colon_eq::Bool,
-    varinfo::AbstractVarInfo=VarInfo(),
+    varinfo::AbstractVarInfo,
     context::AbstractContext=DefaultContext(),
 )
     context = ValuesAsInModelContext(include_colon_eq, context)
     evaluate!!(model, varinfo, context)
     return context.values
 end
-function values_as_in_model(
-    rng::Random.AbstractRNG,
-    model::Model,
-    include_colon_eq::Bool,
-    varinfo::AbstractVarInfo=VarInfo(),
-    context::AbstractContext=DefaultContext(),
-)
-    return values_as_in_model(model, true, varinfo, SamplingContext(rng, context))
-end

src/varinfo.jl

@@ -521,73 +521,22 @@ function merge_metadata(metadata_left::Metadata, metadata_right::Metadata)
     offset = 0
 
     for (idx, vn) in enumerate(vns_both)
-        # `idcs`
         idcs[vn] = idx
-        # `vns`
         push!(vns, vn)
-        if vn in vns_left && vn in vns_right
-            # `vals`: only valid if they're the length.
-            vals_left = getindex_internal(metadata_left, vn)
-            vals_right = getindex_internal(metadata_right, vn)
-            @assert length(vals_left) == length(vals_right)
-            append!(vals, vals_right)
-            # `ranges`
-            r = (offset + 1):(offset + length(vals_left))
-            push!(ranges, r)
-            offset = r[end]
-            # `dists`: only valid if they're the same.
-            dist_right = getdist(metadata_right, vn)
-            # Give precedence to `metadata_right`.
-            push!(dists, dist_right)
-            gid = metadata_right.gids[getidx(metadata_right, vn)]
-            push!(gids, gid)
-            # `orders`: giving precedence to `metadata_right`
-            push!(orders, getorder(metadata_right, vn))
-            # `flags`
-            for k in keys(flags)
-                # Using `metadata_right`; should we?
-                push!(flags[k], is_flagged(metadata_right, vn, k))
-            end
-        elseif vn in vns_left
-            # Just extract the metadata from `metadata_left`.
-            # `vals`
-            vals_left = getindex_internal(metadata_left, vn)
-            append!(vals, vals_left)
-            # `ranges`
-            r = (offset + 1):(offset + length(vals_left))
-            push!(ranges, r)
-            offset = r[end]
-            # `dists`
-            dist_left = getdist(metadata_left, vn)
-            push!(dists, dist_left)
-            gid = metadata_left.gids[getidx(metadata_left, vn)]
-            push!(gids, gid)
-            # `orders`
-            push!(orders, getorder(metadata_left, vn))
-            # `flags`
-            for k in keys(flags)
-                push!(flags[k], is_flagged(metadata_left, vn, k))
-            end
-        else
-            # Just extract the metadata from `metadata_right`.
-            # `vals`
-            vals_right = getindex_internal(metadata_right, vn)
-            append!(vals, vals_right)
-            # `ranges`
-            r = (offset + 1):(offset + length(vals_right))
-            push!(ranges, r)
-            offset = r[end]
-            # `dists`
-            dist_right = getdist(metadata_right, vn)
-            push!(dists, dist_right)
-            gid = metadata_right.gids[getidx(metadata_right, vn)]
-            push!(gids, gid)
-            # `orders`
-            push!(orders, getorder(metadata_right, vn))
-            # `flags`
-            for k in keys(flags)
-                push!(flags[k], is_flagged(metadata_right, vn, k))
-            end
+        metadata_for_vn = vn in vns_right ? metadata_right : metadata_left
+
+        val = getindex_internal(metadata_for_vn, vn)
+        append!(vals, val)
+        r = (offset + 1):(offset + length(val))
+        push!(ranges, r)
+        offset = r[end]
+        dist = getdist(metadata_for_vn, vn)
+        push!(dists, dist)
+        gid = metadata_for_vn.gids[getidx(metadata_for_vn, vn)]
+        push!(gids, gid)
+        push!(orders, getorder(metadata_for_vn, vn))
+        for k in keys(flags)
+            push!(flags[k], is_flagged(metadata_for_vn, vn, k))
         end
     end
 

test/model.jl

@@ -429,22 +429,6 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
                 end
             end
         end
-
-        @testset "check that sampling obeys rng if passed" begin
-            @model function f()
-                x ~ Normal(0)
-                return y ~ Normal(x)
-            end
-            model = f()
-            # Call values_as_in_model with the rng
-            values = values_as_in_model(Random.Xoshiro(43), model, false)
-            # Check that they match the values that would be used if vi was seeded
-            # with that seed instead
-            expected_vi = VarInfo(Random.Xoshiro(43), model)
-            for vn in keys(values)
-                @test values[vn] == expected_vi[vn]
-            end
-        end
     end
 
     @testset "Erroneous model call" begin

test/varinfo.jl

@@ -869,6 +869,17 @@ end
             @test DynamicPPL.getgid(varinfo_merged, @varname(x)) == gidset_left
             @test DynamicPPL.getgid(varinfo_merged, @varname(y)) == gidset_right
         end
+
+        # The below used to error, testing to avoid regression.
+        @testset "merge different dimensions" begin
+            vn = @varname(x)
+            vi_single = VarInfo()
+            vi_single = push!!(vi_single, vn, 1.0, Normal())
+            vi_double = VarInfo()
+            vi_double = push!!(vi_double, vn, [0.5, 0.6], Dirichlet(2, 1.0))
+            @test merge(vi_single, vi_double)[vn] == [0.5, 0.6]
+            @test merge(vi_double, vi_single)[vn] == 1.0
+        end
     end
 
     @testset "VarInfo with selectors" begin