Merge pull request #264 from ReactiveBayes/fix262

Fix issue 262

Author: bvdmitri

src/graph_engine.jl

@@ -341,6 +341,12 @@ end
 getname(label::ProxyLabel) = label.name
 index(label::ProxyLabel) = label.index
 
+# This function allows to overwrite the `maycreate` flag on a proxy label, might be useful for situations where code should
+# definitely not create a new variable, e.g in the variational constraints plugin
+set_maycreate(proxylabel::ProxyLabel, maycreate::Union{True, False}) =
+    ProxyLabel(proxylabel.name, proxylabel.proxied, proxylabel.index, maycreate)
+set_maycreate(something, maycreate::Union{True, False}) = something
+
 function unroll(something)
     return something
 end

src/plugins/variational_constraints/variational_constraints_engine.jl

@@ -754,21 +754,25 @@ function resolve(model::Model, context::Context, variable::IndexedVariable{<:Spl
     )
 end
 
+# The variational constraints plugin should not attempt to create new variables in the model
+# Even if the `maycreate` flag was set to `True` at this point we assume that the variable has been created already
+unroll_nocreate(something) = unroll(set_maycreate(something, False()))
+
 function resolve(model::Model, context::Context, variable::IndexedVariable{Nothing})
-    global_label = unroll(context[getname(variable)])
+    global_label = unroll_nocreate(context[getname(variable)])
     return __resolve(model, global_label)
 end
 
 function resolve(model::Model, context::Context, variable::IndexedVariable)
-    global_label = unroll(context[getname(variable)])[index(variable)...]
+    global_label = unroll_nocreate(context[getname(variable)])[index(variable)...]
     return __resolve(model, global_label)
 end
 
 resolve(model::Model, context::Context, variable::IndexedVariable{CombinedRange{NTuple{N, Int}, NTuple{N, Int}}}) where {N} =
     throw(UnresolvableFactorizationConstraintError("Cannot resolve factorization constraint for a combined range of dimension > 2."))
 
 function resolve(model::Model, context::Context, variable::IndexedVariable{<:CombinedRange})
-    global_label = view(unroll(context[getname(variable)]), firstindex(index(variable)):lastindex(index(variable)))
+    global_label = view(unroll_nocreate(context[getname(variable)]), firstindex(index(variable)):lastindex(index(variable)))
     return __resolve(model, global_label)
 end
 
@@ -803,7 +807,8 @@ function resolve(model::Model, context::Context, constraint::FactorizationConstr
     end
     lhs = map(variable -> resolve(model, context, variable), vfiltered)
     rhs = map(
-        variable -> ResolvedFactorizationConstraintEntry((resolve(model, context, unroll(context[getname(variable)]), MeanField()),)),
+        variable ->
+            ResolvedFactorizationConstraintEntry((resolve(model, context, unroll_nocreate(context[getname(variable)]), MeanField()),)),
         vfiltered
     )
     return ResolvedFactorizationConstraint(ResolvedConstraintLHS(lhs), rhs)
@@ -930,7 +935,7 @@ end
 function apply_constraints!(
     model::Model, context::Context, marginal_constraint::MarginalFormConstraint{T, F} where {T <: IndexedVariable, F}
 )
-    applicable_nodes = unroll(context[getvariables(marginal_constraint)])
+    applicable_nodes = unroll_nocreate(context[getvariables(marginal_constraint)])
     for node in applicable_nodes
         if hasextra(model[node], VariationalConstraintsMarginalFormConstraintKey)
             @warn lazy"Node $node already has functional form constraint $(opt[:q]) applied, therefore $constraint_data will not be applied"
@@ -945,7 +950,7 @@ function apply_constraints!(model::Model, context::Context, marginal_constraint:
 end
 
 function apply_constraints!(model::Model, context::Context, message_constraint::MessageFormConstraint)
-    applicable_nodes = unroll(context[getvariables(message_constraint)])
+    applicable_nodes = unroll_nocreate(context[getvariables(message_constraint)])
     for node in applicable_nodes
         if hasextra(model[node], VariationalConstraintsMessagesFormConstraintKey)
             @warn lazy"Node $node already has functional form constraint $(opt[:q]) applied, therefore $constraint_data will not be applied"

test/graph_engine_tests.jl

@@ -712,6 +712,7 @@ end
         getcontext,
         getifcreated,
         unroll,
+        set_maycreate,
         ProxyLabel,
         NodeLabel,
         proxylabel,
@@ -812,6 +813,27 @@ end
         # `x` can be inferred properly
         @test ctx[:x] === @inferred(unroll(proxylabel(:x, xref, nothing, False())))
     end
+
+    @testset "It should be possible to toggle `maycreate` flag" begin
+        model = create_test_model()
+        ctx = getcontext(model)
+        xref = VariableRef(model, ctx, NodeCreationOptions(), :x, (nothing,))
+        # The first time should throw since the variable has not been instantiated yet
+        @test_throws "The variable `x` has been used, but has not been instantiated." unroll(proxylabel(:x, xref, nothing, False()))
+        # Even though the `maycreate` flag is set to `True`, the `set_maycreate` should overwrite it with `False`
+        @test_throws "The variable `x` has been used, but has not been instantiated." unroll(
+            set_maycreate(proxylabel(:x, xref, nothing, True()), False())
+        )
+
+        # Even though the `maycreate` flag is set to `False`, the `set_maycreate` should overwrite it with `True`
+        @test unroll(set_maycreate(proxylabel(:x, xref, nothing, False()), True())) === ctx[:x]
+        # At this point the variable should be created
+        @test unroll(proxylabel(:x, xref, nothing, False())) === ctx[:x]
+        @test unroll(proxylabel(:x, xref, nothing, True())) === ctx[:x]
+
+        @test set_maycreate(1, True()) === 1
+        @test set_maycreate(1, False()) === 1
+    end
 end
 
 @testitem "`VariableRef` comparison" begin

test/plugins/variational_constraints/variational_constraints_tests.jl

@@ -1107,3 +1107,30 @@ end
         end
     end
 end
+
+@testitem "Issue 262, factorization constraint should not attempt to create a variable from submodels" begin
+    import GraphPPL: create_model, with_plugins, getproperties, neighbor_data, is_factorized
+
+    include("../../testutils.jl")
+
+    @model function submodel(y, x)
+        for i in 1:10
+            y[i] ~ Normal(x, 1)
+        end
+    end
+
+    @model function main_model()
+        x ~ Normal(0, 1)
+        y ~ submodel(x = x)
+    end
+
+    constraints = @constraints begin
+        for q in submodel
+            q(x, y) = q(x)q(y)
+        end
+    end
+
+    model = create_model(with_plugins(main_model(), GraphPPL.PluginsCollection(GraphPPL.VariationalConstraintsPlugin(constraints))))
+
+    @test length(collect(filter(as_node(Normal), model))) == 11
+end