save indices in the extra field

ReactiveBayes · Feb 26, 2024 · 30bc233 · 30bc233
1 parent 17cb6fd
commit 30bc233
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/plugins/variational_constraints/variational_constraints_engine.jl b/src/plugins/variational_constraints/variational_constraints_engine.jl
@@ -588,12 +588,13 @@ function materialize_constraints!(model::Model, node_label::NodeLabel, node_data
         )
     end
 
-    # TODO: (bvdmitri) use node properties for faster access
+    # TODO: (bvdmitri) use node properties for faster access (we dont even need this, right?)
     edges = GraphPPL.edges(model, node_label)
     new_constraint = Tuple(sort!(collect(constraint_set), by = first))
     new_constraint_ = map(clusters -> Tuple(getindex.(Ref(edges), clusters)), new_constraint)
 
     setextra!(node_data, :factorization_constraint, new_constraint_)
+    setextra!(node_data, :factorization_constraint_indices, new_constraint)
 end
 
 function materialize_constraints!(model::Model, node_label::NodeLabel, node_data::NodeData, ::VariableNodeProperties)