Merge pull request #302 from ReactiveBayes/inline_constraints

Add support to specify constraints macros inline for submodels

Author: bvdmitri

docs/src/plugins/constraint_specification.md

@@ -78,6 +78,42 @@ We can specify constraints over the first `toy_model` submodel using the followi
 end
 ```
 
+## Inline constraints on submodel calls
+
+Constraints can also be specified **inline** at the call site of a submodel using the `where` keyword. This is useful when you want to attach constraints to a specific invocation without modifying the outer model's `@constraints` block. For example:
+
+```@example constraints
+@model function outer_toy_model(a, b, c)
+    a ~ toy_model(y = b, z = c) where {
+        constraints = @constraints begin
+            q(x, y, z) = q(x, y)q(z)
+            q(x) :: Normal
+        end
+    }
+end
+```
+
+The `where { constraints = ... }` syntax accepts any constraint set produced by the `@constraints` macro. Constraint sets defined with the `@constraints function` form can also be passed by reference:
+
+```@example constraints
+@constraints function my_constraints()
+    q(x, y, z) = q(x, y)q(z)
+    q(x) :: Normal
+end
+
+@model function outer_toy_model(a, b, c)
+    a ~ toy_model(y = b, z = c) where { constraints = my_constraints() }
+end
+```
+
+Inline constraints apply only to the specific submodel invocation they are attached to and propagate to any submodels nested within it. Their priority relative to other constraint sources, from highest to lowest, is:
+
+1. **External constraints** — passed at model creation via `for q in submodel` or `for q in (submodel, index)` blocks.
+2. **Inline constraints** — specified with `where { constraints = ... }` at the call site.
+3. **Default constraints** — defined via `GraphPPL.default_constraints`.
+
+This means that if external constraints also target the same submodel, they will override the inline constraints.
+
 ## Constraints over vector variables
 
 When a model contains vector (or array) latent variables, we can specify factorization constraints over individual elements using the `begin` and `end` indexing syntax. For example, consider a random walk model where latent states `x` are coupled through sequential dependencies:

src/graph_engine.jl

@@ -452,9 +452,10 @@ struct Context
     tensor_variables::UnorderedDictionary{Symbol, ResizableArray{NodeLabel}}
     proxies::UnorderedDictionary{Symbol, ProxyLabel}
     returnval::Ref{Any}
+    options::Any
 end
 
-function Context(depth::Int, fform::Function, prefix::String, parent)
+function Context(depth::Int, fform::Function, prefix::String, parent, options = nothing)
     return Context(
         depth,
         fform,
@@ -467,13 +468,17 @@ function Context(depth::Int, fform::Function, prefix::String, parent)
         UnorderedDictionary{Symbol, ResizableArray{NodeLabel, Vector{NodeLabel}, 1}}(),
         UnorderedDictionary{Symbol, ResizableArray{NodeLabel}}(),
         UnorderedDictionary{Symbol, ProxyLabel}(),
-        Ref{Any}()
+        Ref{Any}(),
+        options
     )
 end
 
 Context(parent::Context, model_fform::Function) = Context(
     parent.depth + 1, model_fform, (parent.prefix == "" ? parent.prefix : parent.prefix * "_") * getname(model_fform), parent
 )
+Context(parent::Context, model_fform::Function, options) = Context(
+    parent.depth + 1, model_fform, (parent.prefix == "" ? parent.prefix : parent.prefix * "_") * getname(model_fform), parent, options
+)
 Context(fform) = Context(0, fform, "", nothing)
 Context() = Context(identity)
 
@@ -660,6 +665,8 @@ Base.getindex(::NodeCreationOptions{Nothing}, keys...) = error("type `NodeCreati
 Base.keys(::NodeCreationOptions{Nothing}) = ()
 Base.get(::NodeCreationOptions{Nothing}, key::Symbol, default) = default
 
+context_options(context::Context) = something(context.options, EmptyNodeCreationOptions)
+
 withopts(::NodeCreationOptions{Nothing}, options::NamedTuple) = NodeCreationOptions(options)
 withopts(options::NodeCreationOptions, extra::NamedTuple) = NodeCreationOptions((; options.options..., extra...))
 

src/model_macro.jl

@@ -771,7 +771,7 @@ function get_make_node_function(model_specification, ms_body, ms_args, ms_name)
             __n_interfaces__::GraphPPL.StaticInt{$(length(ms_args))}
         )
             __interfaces__ = GraphPPL.prepare_interfaces(__model__, $ms_name, __lhs_interface__, __rhs_interfaces__)
-            __context__ = GraphPPL.Context(__parent_context__, $ms_name)
+            __context__ = GraphPPL.Context(__parent_context__, $ms_name, __options__)
             GraphPPL.copy_markov_blanket_to_child_context(__context__, __interfaces__)
             GraphPPL.add_composite_factor_node!(__model__, __parent_context__, __context__, $ms_name)
             __returnval__ = GraphPPL.add_terminated_submodel!(

src/plugins/variational_constraints/variational_constraints_engine.jl

@@ -1000,10 +1000,13 @@ function apply_constraints!(model::Model, context::Context, constraint_set::Cons
         apply_constraints!(model, context, rfc)
     end
     for (factor_id, child) in pairs(children(context))
+        inline_constraints = get(context_options(child), :constraints, nothing)
         if factor_id ∈ keys(specific_submodel_constraints(constraint_set))
             apply_constraints!(model, child, getconstraint(specific_submodel_constraints(constraint_set)[factor_id]), stack)
         elseif fform(factor_id) ∈ keys(general_submodel_constraints(constraint_set))
             apply_constraints!(model, child, getconstraint(general_submodel_constraints(constraint_set)[fform(child)]), stack)
+        elseif inline_constraints !== nothing
+            apply_constraints!(model, child, inline_constraints, stack)
         else
             apply_constraints!(model, child, default_constraints(fform(factor_id)), stack)
         end

test/plugins/variational_constraints/variational_constraints_tests.jl

@@ -1132,6 +1132,92 @@ end
     @test length(collect(filter(as_node(Normal), model))) == 11
 end
 
+@testitem "inline constraints on submodel calls" begin
+    using Distributions
+    import GraphPPL:
+        create_model,
+        with_plugins,
+        getcontext,
+        getextra,
+        hasextra,
+        context_options,
+        VariationalConstraintsPlugin,
+        PluginsCollection,
+        Constraints
+
+    include("../../testutils.jl")
+
+    @model function inline_constraints_inner(y, x)
+        theta ~ Normal(0.0, 1.0)
+        y ~ Normal(x, theta)
+    end
+
+    @model function inline_constraints_middle(y, x)
+        y ~ inline_constraints_inner(x = x) where {
+            constraints = @constraints begin
+                q(x, y, theta) = MeanField()
+            end
+        }
+    end
+
+    @model function inline_constraints_outer()
+        x ~ Normal(0.0, 1.0)
+        y ~ inline_constraints_inner(x = x) where {
+            constraints = @constraints begin
+                q(x, y, theta) = MeanField()
+            end
+        }
+    end
+
+    @constraints function make_inline_constraints()
+        q(x, y, theta) = MeanField()
+    end
+
+    @model function inline_constraints_function_outer()
+        x ~ Normal(0.0, 1.0)
+        y ~ inline_constraints_inner(x = x) where { constraints = make_inline_constraints() }
+    end
+
+    @model function nested_inline_constraints_outer()
+        x ~ Normal(0.0, 1.0)
+        y ~ inline_constraints_middle(x = x)
+    end
+
+    model = create_model(with_plugins(inline_constraints_outer(), PluginsCollection(VariationalConstraintsPlugin())))
+    context = getcontext(model)
+    inner_context = context[inline_constraints_inner, 1]
+    node = inner_context[NormalMeanVariance, 2]
+    @test get(context_options(inner_context), :constraints, nothing) isa Constraints
+    @test hasextra(model[node], :factorization_constraint_indices)
+    @test Tuple.(getextra(model[node], :factorization_constraint_indices)) == ((1,), (2,), (3,))
+
+    model = create_model(with_plugins(inline_constraints_function_outer(), PluginsCollection(VariationalConstraintsPlugin())))
+    context = getcontext(model)
+    inner_context = context[inline_constraints_inner, 1]
+    node = inner_context[NormalMeanVariance, 2]
+    @test get(context_options(inner_context), :constraints, nothing) isa Constraints
+    @test hasextra(model[node], :factorization_constraint_indices)
+    @test Tuple.(getextra(model[node], :factorization_constraint_indices)) == ((1,), (2,), (3,))
+
+    model = create_model(with_plugins(nested_inline_constraints_outer(), PluginsCollection(VariationalConstraintsPlugin())))
+    context = getcontext(model)
+    inner_context = context[inline_constraints_middle, 1][inline_constraints_inner, 1]
+    node = inner_context[NormalMeanVariance, 2]
+    @test hasextra(model[node], :factorization_constraint_indices)
+    @test Tuple.(getextra(model[node], :factorization_constraint_indices)) == ((1,), (2,), (3,))
+
+    external_constraints = @constraints begin
+        for q in inline_constraints_inner
+            q(x, y, theta) = q(x, y, theta)
+        end
+    end
+    model = create_model(with_plugins(inline_constraints_outer(), PluginsCollection(VariationalConstraintsPlugin(external_constraints))))
+    context = getcontext(model)
+    node = context[inline_constraints_inner, 1][NormalMeanVariance, 2]
+    @test hasextra(model[node], :factorization_constraint_indices)
+    @test Tuple.(getextra(model[node], :factorization_constraint_indices)) == ((1, 2, 3),)
+end
+
 @testitem "`@constraints` should save the source code #1" begin
     using GraphPPL