Implement tracked_varnames

penelopeysm · penelopeysm · commit bb8adb77af48 · 2025-03-18T01:05:54.000Z
diff --git a/src/values_as_in_model.jl b/src/values_as_in_model.jl
@@ -1,3 +1,12 @@
+"""
+    TrackedValue{T}
+
+A struct that wraps something on the right-hand side of `:=`. This is needed
+because the DynamicPPL compiler actually converts `lhs := rhs` to `lhs ~
+TrackedValue(rhs)` (so that we can hit the `tilde_assume` method below). Having
+the rhs wrapped in a TrackedValue makes sure that the logpdf of the rhs is not
+computed (as it wouldn't make sense).
+"""
 struct TrackedValue{T}
     value::T
 end
@@ -24,17 +33,27 @@ struct ValuesAsInModelContext{C<:AbstractContext} <: AbstractContext
     values::OrderedDict
     "whether to extract variables on the LHS of :="
     include_colon_eq::Bool
+    "varnames to be tracked; `nothing` means track all varnames"
+    tracked_varnames::Union{Nothing,Array{<:VarName}}
     "child context"
     context::C
 end
-function ValuesAsInModelContext(include_colon_eq, context::AbstractContext)
-    return ValuesAsInModelContext(OrderedDict(), include_colon_eq, context)
+function ValuesAsInModelContext(
+    include_colon_eq::Bool,
+    tracked_varnames::Union{Nothing,Array{<:VarName}},
+    context::AbstractContext,
+)
+    return ValuesAsInModelContext(
+        OrderedDict(), include_colon_eq, tracked_varnames, context
+    )
 end
 
 NodeTrait(::ValuesAsInModelContext) = IsParent()
 childcontext(context::ValuesAsInModelContext) = context.context
 function setchildcontext(context::ValuesAsInModelContext, child)
-    return ValuesAsInModelContext(context.values, context.include_colon_eq, child)
+    return ValuesAsInModelContext(
+        context.values, context.include_colon_eq, context.tracked_varnames, child
+    )
 end
 
 is_extracting_values(context::ValuesAsInModelContext) = context.include_colon_eq
@@ -63,29 +82,38 @@ end
 
 # `tilde_asssume`
 function tilde_assume(context::ValuesAsInModelContext, right, vn, vi)
-    if is_tracked_value(right)
+    is_tracked_value_right = is_tracked_value(right)
+    if is_tracked_value_right
         value = right.value
         logp = zero(getlogp(vi))
     else
         value, logp, vi = tilde_assume(childcontext(context), right, vn, vi)
     end
     # Save the value.
-    push!(context, vn, value)
-    # Save the value.
+    if is_tracked_value_right ||
+        isnothing(context.tracked_varnames) ||
+        any(tracked_vn -> subsumes(tracked_vn, vn), context.tracked_varnames)
+        push!(context, vn, value)
+    end
     # Pass on.
     return value, logp, vi
 end
 function tilde_assume(
     rng::Random.AbstractRNG, context::ValuesAsInModelContext, sampler, right, vn, vi
 )
-    if is_tracked_value(right)
+    is_tracked_value_right = is_tracked_value(right)
+    if is_tracked_value_right
         value = right.value
         logp = zero(getlogp(vi))
     else
         value, logp, vi = tilde_assume(rng, childcontext(context), sampler, right, vn, vi)
     end
     # Save the value.
-    push!(context, vn, value)
+    if is_tracked_value_right ||
+        isnothing(context.tracked_varnames) ||
+        any(tracked_vn -> subsumes(tracked_vn, vn), context.tracked_varnames)
+        push!(context, vn, value)
+    end
     # Pass on.
     return value, logp, vi
 end
@@ -167,9 +195,39 @@ function values_as_in_model(
     model::Model,
     include_colon_eq::Bool,
     varinfo::AbstractVarInfo,
+    tracked_varnames=tracked_varnames(model),
     context::AbstractContext=DefaultContext(),
 )
-    context = ValuesAsInModelContext(include_colon_eq, context)
+    tracked_varnames = isnothing(tracked_varnames) ? nothing : collect(tracked_varnames)
+    context = ValuesAsInModelContext(include_colon_eq, tracked_varnames, context)
     evaluate!!(model, varinfo, context)
     return context.values
 end
+
+"""
+    tracked_varnames(model::Model)
+
+Returns a set of `VarName`s that the model should track.
+
+By default, this returns `nothing`, which means that all `VarName`s should be
+tracked.
+
+If you want to track only a subset of `VarName`s, you can override this method
+in your model definition:
+
+```julia
+@model function mymodel()
+    x ~ Normal()
+    y ~ Normal(x, 1)
+end
+
+DynamicPPL.tracked_varnames(::Model{typeof(mymodel)}) = [@varname(y)]
+```
+
+Then, when you sample from `mymodel()`, only the value of `y` will be tracked
+(and not `x`).
+
+Note that quantities on the left-hand side of `:=` are always tracked, and will
+ignore the varnames specified in this method.
+"""
+tracked_varnames(::Model) = nothing
diff --git a/test/compiler.jl b/test/compiler.jl
@@ -728,18 +728,6 @@ module Issue537 end
             varinfo = VarInfo(model)
             @test haskey(varinfo, @varname(x))
             @test !haskey(varinfo, @varname(y))
-
-            # While `values_as_in_model` should contain both `x` and `y`, if
-            # include_colon_eq is set to `true`.
-            values = values_as_in_model(model, true, deepcopy(varinfo))
-            @test haskey(values, @varname(x))
-            @test haskey(values, @varname(y))
-
-            # And if include_colon_eq is set to `false`, then `values` should
-            # only contain `x`.
-            values = values_as_in_model(model, false, deepcopy(varinfo))
-            @test haskey(values, @varname(x))
-            @test !haskey(values, @varname(y))
         end
     end
 
diff --git a/test/model.jl b/test/model.jl
@@ -410,48 +410,6 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
         end
     end
 
-    @testset "values_as_in_model" begin
-        @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
-            vns = DynamicPPL.TestUtils.varnames(model)
-            example_values = DynamicPPL.TestUtils.rand_prior_true(model)
-            varinfos = DynamicPPL.TestUtils.setup_varinfos(model, example_values, vns)
-            @testset "$(short_varinfo_name(varinfo))" for varinfo in varinfos
-                # We can set the include_colon_eq arg to false because none of
-                # the demo models contain :=. The behaviour when
-                # include_colon_eq is true is tested in test/compiler.jl
-                realizations = values_as_in_model(model, false, varinfo)
-                # Ensure that all variables are found.
-                vns_found = collect(keys(realizations))
-                @test vns ∩ vns_found == vns ∪ vns_found
-                # Ensure that the values are the same.
-                for vn in vns
-                    @test realizations[vn] == varinfo[vn]
-                end
-            end
-        end
-
-        @testset "Prefixing" begin
-            @model inner() = x ~ Normal()
-
-            @model function outer_auto_prefix()
-                a ~ to_submodel(inner(), true)
-                b ~ to_submodel(inner(), true)
-                return nothing
-            end
-            @model function outer_manual_prefix()
-                a ~ to_submodel(prefix(inner(), :a), false)
-                b ~ to_submodel(prefix(inner(), :b), false)
-                return nothing
-            end
-
-            for model in (outer_auto_prefix(), outer_manual_prefix())
-                vi = VarInfo(model)
-                vns = Set(keys(values_as_in_model(model, false, vi)))
-                @test vns == Set([@varname(var"a.x"), @varname(var"b.x")])
-            end
-        end
-    end
-
     @testset "Erroneous model call" begin
         # Calling a model with the wrong arguments used to lead to infinite recursion, see
         # https://github.com/TuringLang/Turing.jl/issues/2182. This guards against it.
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -51,6 +51,7 @@ include("test_util.jl")
         include("varinfo.jl")
         include("simple_varinfo.jl")
         include("model.jl")
+        include("values_as_in_model.jl")
         include("sampler.jl")
         include("independence.jl")
         include("distribution_wrappers.jl")
diff --git a/test/values_as_in_model.jl b/test/values_as_in_model.jl
@@ -0,0 +1,93 @@
+@testset "values_as_in_model" begin
+    @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
+        vns = DynamicPPL.TestUtils.varnames(model)
+        example_values = DynamicPPL.TestUtils.rand_prior_true(model)
+        varinfos = DynamicPPL.TestUtils.setup_varinfos(model, example_values, vns)
+        @testset "$(short_varinfo_name(varinfo))" for varinfo in varinfos
+            # We can set the include_colon_eq arg to false because none of
+            # the demo models contain :=. The behaviour when
+            # include_colon_eq is true is tested in test/compiler.jl
+            realizations = values_as_in_model(model, false, varinfo)
+            # Ensure that all variables are found.
+            vns_found = collect(keys(realizations))
+            @test vns ∩ vns_found == vns ∪ vns_found
+            # Ensure that the values are the same.
+            for vn in vns
+                @test realizations[vn] == varinfo[vn]
+            end
+        end
+    end
+
+    @testset "support for :=" begin
+        @model function demo_tracked()
+            x ~ Normal()
+            y := 100 + x
+            return (; x, y)
+        end
+        @model function demo_tracked_submodel()
+            return vals ~ to_submodel(demo_tracked(), false)
+        end
+
+        for model in [demo_tracked(), demo_tracked_submodel()]
+            values = values_as_in_model(model, true, VarInfo(model))
+            @test haskey(values, @varname(x))
+            @test haskey(values, @varname(y))
+
+            values = values_as_in_model(model, false, VarInfo(model))
+            @test haskey(values, @varname(x))
+            @test !haskey(values, @varname(y))
+        end
+    end
+
+    @testset "Prefixing" begin
+        @model inner() = x ~ Normal()
+
+        @model function outer_auto_prefix()
+            a ~ to_submodel(inner(), true)
+            b ~ to_submodel(inner(), true)
+            return nothing
+        end
+        @model function outer_manual_prefix()
+            a ~ to_submodel(prefix(inner(), :a), false)
+            b ~ to_submodel(prefix(inner(), :b), false)
+            return nothing
+        end
+
+        for model in (outer_auto_prefix(), outer_manual_prefix())
+            vi = VarInfo(model)
+            vns = Set(keys(values_as_in_model(model, false, vi)))
+            @test vns == Set([@varname(var"a.x"), @varname(var"b.x")])
+        end
+    end
+
+    @testset "Track only specific varnames" begin
+        @model function track_specific()
+            x = Vector{Float64}(undef, 2)
+            # Include a vector x to test for correct subsumption behaviour
+            for i in eachindex(x)
+                x[i] ~ Normal()
+            end
+            y ~ Normal(x[1], 1)
+            z := sum(x)
+        end
+
+        model = track_specific()
+        vi = VarInfo(model)
+
+        # Specify varnames to be tracked directly as an argument to `values_as_in_model`
+        values = values_as_in_model(model, true, vi, [@varname(x)])
+        # Since x subsumes both x[1] and x[2], they should be included
+        @test haskey(values, @varname(x[1]))
+        @test haskey(values, @varname(x[2]))
+        @test !haskey(values, @varname(y))
+        @test haskey(values, @varname(z))  # := is always included
+
+        # Specify instead using `tracked_varnames` method
+        DynamicPPL.tracked_varnames(::Model{typeof(track_specific)}) = [@varname(y)]
+        values = values_as_in_model(model, true, vi)
+        @test !haskey(values, @varname(x[1]))
+        @test !haskey(values, @varname(x[2]))
+        @test haskey(values, @varname(y))
+        @test haskey(values, @varname(z))  # := is always included
+    end
+end
