Skip to content

Commit 08a1973

Browse files
bvdmitribvdmitri
authored
Merge pull request #301 from ReactiveBayes/multioutput_submodels
Multi-output submodels
2 parents 07e31d0 + 3687284 commit 08a1973

5 files changed

Lines changed: 487 additions & 13 deletions

File tree

‎docs/src/nested_models.md‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,4 +30,43 @@ If we want to chain these `gcv` submodels together into a Hierarchical Gaussian
3030
end
3131
```
3232

33-
Note that in our invocations of `gcv`, we haven't specified the `y` argument of the Markov Blanket. This is what is being recognized as the missing interface and `GraphPPL` will assign `means[i]` to `y`.
33+
Note that in our invocations of `gcv`, we haven't specified the `y` argument of the Markov Blanket. This is what is being recognized as the missing interface and `GraphPPL` will assign `means[i]` to `y`.
34+
35+
## Multi-output submodels
36+
37+
When a submodel produces multiple outputs — multiple interfaces left unspecified on the RHS — you can bind them all on the LHS using a tuple. There are two syntaxes:
38+
39+
**Positional:** list outer variables in the same order as the unspecified interfaces appear in the submodel definition.
40+
41+
``` @example nested-models
42+
@model function linear_gaussian(x, y, z)
43+
x ~ Normal(z, 1)
44+
y ~ Normal(x, 1)
45+
end
46+
47+
@model function outer_positional(c)
48+
(a, b) ~ linear_gaussian(z = c) # a → interface x, b → interface y (by position)
49+
end
50+
```
51+
52+
**Named (kwarg-style):** explicitly map each outer variable to its interface name using `name = var` pairs. This is order-independent and recommended when submodel argument order may change.
53+
54+
``` @example nested-models
55+
@model function outer_named(my_z)
56+
(y = my_y, x = my_x) ~ linear_gaussian(z = my_z) # binds by name, regardless of order
57+
obs ~ Normal(my_x, my_y)
58+
end
59+
```
60+
61+
Both syntaxes work with indexed variables in loops:
62+
63+
``` @example nested-models
64+
@model function chain(z, n)
65+
for i in 1:n
66+
(x = xs[i], y = ys[i]) ~ linear_gaussian(z = z)
67+
end
68+
end
69+
```
70+
71+
!!! note
72+
If the same interface name appears on both LHS and RHS, `GraphPPL` raises an error at model-creation time. Similarly, providing a name on the LHS that does not match any of the submodel's interface names is caught with a descriptive error.

‎src/graph_engine.jl‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1203,6 +1203,7 @@ struct StaticInterfaces{I} end
12031203

12041204
StaticInterfaces(I::Tuple) = StaticInterfaces{I}()
12051205
Base.getindex(::StaticInterfaces{I}, index) where {I} = I[index]
1206+
iface_names(::StaticInterfaces{I}) where {I} = I
12061207

12071208
function Base.convert(::Type{NamedTuple}, ::StaticInterfaces{I}, t::Tuple) where {I}
12081209
return NamedTuple{I}(t)
@@ -1870,6 +1871,37 @@ function prepare_interfaces(::StaticInterfaces{I}, fform::F, lhs_interface, rhs_
18701871
return NamedTuple{(missing_interface, keys(rhs_interfaces)...)}((lhs_interface, values(rhs_interfaces)...))
18711872
end
18721873

1874+
# Multi-output: lhs_interfaces is a Tuple of multiple interfaces (positional)
1875+
function prepare_interfaces(model::Model, fform::F, lhs_interfaces::Tuple, rhs_interfaces::NamedTuple) where {F}
1876+
n_lhs = length(lhs_interfaces)
1877+
missing = missing_interfaces(model, fform, static(length(rhs_interfaces) + n_lhs), rhs_interfaces)
1878+
return prepare_interfaces_multi(missing, fform, lhs_interfaces, rhs_interfaces)
1879+
end
1880+
1881+
function prepare_interfaces_multi(::StaticInterfaces{I}, fform::F, lhs_interfaces::Tuple, rhs_interfaces::NamedTuple) where {I, F}
1882+
all_keys = (I..., keys(rhs_interfaces)...)
1883+
all_vals = (lhs_interfaces..., values(rhs_interfaces)...)
1884+
return NamedTuple{all_keys}(all_vals)
1885+
end
1886+
1887+
# Named-output: lhs_interfaces is a NamedTuple (kwarg-style, e.g. (a = m_a, b = m_b) ~ sub(x = x))
1888+
function prepare_interfaces(model::Model, fform::F, lhs_interfaces::NamedTuple, rhs_interfaces::NamedTuple) where {F}
1889+
for k in keys(lhs_interfaces)
1890+
if k ∈ keys(rhs_interfaces)
1891+
error(lazy"Interface ':$(k)' of '$(fform)' is specified on both LHS and RHS.")
1892+
end
1893+
end
1894+
merged = merge(lhs_interfaces, rhs_interfaces)
1895+
all_ifaces = interfaces(model, fform, static(length(merged)))
1896+
valid_names = iface_names(all_ifaces)
1897+
for k in keys(lhs_interfaces)
1898+
if k ∉ valid_names
1899+
error(lazy"Interface ':$(k)' does not exist in '$(fform)'. Valid interfaces are: $(valid_names).")
1900+
end
1901+
end
1902+
return merged
1903+
end
1904+
18731905
function materialize_interface(model, context, interface)
18741906
return getifcreated(model, context, unroll(interface))
18751907
end
@@ -2019,6 +2051,16 @@ make_node!(materialize::True, node_type::Composite, behaviour::Stochastic, model
20192051
Composite(), model, ctx, options, fform, lhs_interface, rhs_interfaces, static(length(rhs_interfaces) + 1)
20202052
)
20212053

2054+
# Multi-output: Tuple LHS for composite nodes (positional)
2055+
make_node!(materialize::True, node_type::Composite, behaviour::Stochastic, model::Model, ctx::Context, options::NodeCreationOptions, fform::F, lhs_interface::Tuple, rhs_interfaces::NamedTuple) where {F} = make_node!(
2056+
Composite(), model, ctx, options, fform, lhs_interface, rhs_interfaces, static(length(rhs_interfaces) + length(lhs_interface))
2057+
)
2058+
2059+
# Named-output: NamedTuple LHS for composite nodes (kwarg-style)
2060+
make_node!(materialize::True, node_type::Composite, behaviour::Stochastic, model::Model, ctx::Context, options::NodeCreationOptions, fform::F, lhs_interface::NamedTuple, rhs_interfaces::NamedTuple) where {F} = make_node!(
2061+
Composite(), model, ctx, options, fform, lhs_interface, rhs_interfaces, static(length(rhs_interfaces) + length(lhs_interface))
2062+
)
2063+
20222064
"""
20232065
make_node!
20242066

‎src/model_macro.jl‎

Lines changed: 119 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -399,15 +399,41 @@ A `quote` block with the modified expression.
399399
"""
400400
function add_get_or_create_expression(e::Expr)
401401
if @capture(e, (lhs_ ~ rhs_ where {options__}))
402-
@capture(lhs, (var_[index__]) | (var_))
403-
return quote
404-
$(generate_get_or_create(var, index, rhs))
405-
$e
402+
if lhs isa Expr && lhs.head == :tuple
403+
if is_named_tuple_lhs(lhs)
404+
# Named-output: (a = m_a, b = m_b) ~ sub(...) — get_or_create the outer (RHS) variables
405+
creates = map(lhs.args) do elem
406+
outer_var_expr = elem.args[2]
407+
@capture(outer_var_expr, (var_[index__]) | (var_)) ||
408+
error("Invalid named LHS value in $(elem). Value must be in a `var` or `var[index]` form.")
409+
generate_get_or_create(var, index, rhs)
410+
end
411+
else
412+
# Positional-output: (m_a, m_b) ~ sub(...)
413+
creates = map(lhs.args) do elem
414+
@capture(elem, (var_[index__]) | (var_)) ||
415+
error("Invalid tuple element on LHS: $(elem). Must be in a `var` or `var[index]` form.")
416+
generate_get_or_create(var, index, rhs)
417+
end
418+
end
419+
return quote
420+
$(creates...)
421+
$e
422+
end
423+
else
424+
@capture(lhs, (var_[index__]) | (var_))
425+
return quote
426+
$(generate_get_or_create(var, index, rhs))
427+
$e
428+
end
406429
end
407430
end
408431
return e
409432
end
410433

434+
is_named_tuple_lhs(lhs) = false
435+
is_named_tuple_lhs(lhs::Expr) = lhs.head === :tuple && !isempty(lhs.args) && all(elem -> elem isa Expr && elem.head === :(=), lhs.args)
436+
411437
what_walk(::typeof(add_get_or_create_expression)) = not_created_by
412438

413439
"""
@@ -647,13 +673,51 @@ function convert_tilde_expression(e::Expr)
647673
options = GraphPPL.options_vector_to_named_tuple(options)
648674
nodesym = gensym(:node)
649675
varsym = gensym(:var)
650-
@capture(lhs, (var_[index__]) | (var_)) || error("Invalid left-hand side $(lhs). Must be in a `var` or `var[index]` form.")
651-
return quote
652-
begin
653-
$nodesym, $varsym = GraphPPL.make_node!(
654-
__model__, __context__, GraphPPL.NodeCreationOptions($(options)), $fform, $(generate_lhs_proxylabel(var, index)), $args
655-
)
656-
$varsym
676+
if lhs isa Expr && lhs.head == :tuple
677+
if is_named_tuple_lhs(lhs)
678+
# Named-output: (a = m_a, b = m_b) ~ sub(...) — build a NamedTuple of proxy labels
679+
proxy_pairs = map(lhs.args) do elem
680+
iface_name = elem.args[1]
681+
outer_var_expr = elem.args[2]
682+
@capture(outer_var_expr, (var_[index__]) | (var_)) ||
683+
error("Invalid named LHS value in $(elem). Value must be in a `var` or `var[index]` form.")
684+
Expr(:(=), iface_name, generate_lhs_proxylabel(var, index))
685+
end
686+
lhs_named = Expr(:tuple, proxy_pairs...)
687+
return quote
688+
begin
689+
$nodesym, $varsym = GraphPPL.make_node!(
690+
__model__, __context__, GraphPPL.NodeCreationOptions($(options)), $fform, $lhs_named, $args
691+
)
692+
$varsym
693+
end
694+
end
695+
else
696+
# Positional-output: (m_a, m_b) ~ sub(...) — build a plain Tuple of proxy labels
697+
proxy_labels = map(lhs.args) do elem
698+
@capture(elem, (var_[index__]) | (var_)) ||
699+
error("Invalid tuple element on LHS: $(elem). Must be in a `var` or `var[index]` form.")
700+
generate_lhs_proxylabel(var, index)
701+
end
702+
lhs_tuple = Expr(:tuple, proxy_labels...)
703+
return quote
704+
begin
705+
$nodesym, $varsym = GraphPPL.make_node!(
706+
__model__, __context__, GraphPPL.NodeCreationOptions($(options)), $fform, $lhs_tuple, $args
707+
)
708+
$varsym
709+
end
710+
end
711+
end
712+
else
713+
@capture(lhs, (var_[index__]) | (var_)) || error("Invalid left-hand side $(lhs). Must be in a `var` or `var[index]` form.")
714+
return quote
715+
begin
716+
$nodesym, $varsym = GraphPPL.make_node!(
717+
__model__, __context__, GraphPPL.NodeCreationOptions($(options)), $fform, $(generate_lhs_proxylabel(var, index)), $args
718+
)
719+
$varsym
720+
end
657721
end
658722
end
659723
elseif @capture(e, (lhs_ .~ fform_(args__; kwargs__) where {options__}) | (lhs_ .~ fform_(args__) where {options__}))
@@ -781,6 +845,50 @@ function get_make_node_function(model_specification, ms_body, ms_args, ms_name)
781845
return __context__, __lhs_interface__
782846
end
783847

848+
# Multi-output: Tuple LHS (positional)
849+
function GraphPPL.make_node!(
850+
::GraphPPL.Composite,
851+
__model__::GraphPPL.Model,
852+
__parent_context__::GraphPPL.Context,
853+
__options__::GraphPPL.NodeCreationOptions,
854+
::typeof($ms_name),
855+
__lhs_interface__::Tuple,
856+
__rhs_interfaces__::NamedTuple,
857+
__n_interfaces__::GraphPPL.StaticInt{$(length(ms_args))}
858+
)
859+
__interfaces__ = GraphPPL.prepare_interfaces(__model__, $ms_name, __lhs_interface__, __rhs_interfaces__)
860+
__context__ = GraphPPL.Context(__parent_context__, $ms_name)
861+
GraphPPL.copy_markov_blanket_to_child_context(__context__, __interfaces__)
862+
GraphPPL.add_composite_factor_node!(__model__, __parent_context__, __context__, $ms_name)
863+
__returnval__ = GraphPPL.add_terminated_submodel!(
864+
__model__, __context__, __options__, $ms_name, __interfaces__, __n_interfaces__
865+
)
866+
GraphPPL.returnval!(__context__, __returnval__)
867+
return __context__, __lhs_interface__
868+
end
869+
870+
# Named-output: NamedTuple LHS (kwarg-style)
871+
function GraphPPL.make_node!(
872+
::GraphPPL.Composite,
873+
__model__::GraphPPL.Model,
874+
__parent_context__::GraphPPL.Context,
875+
__options__::GraphPPL.NodeCreationOptions,
876+
::typeof($ms_name),
877+
__lhs_interface__::NamedTuple,
878+
__rhs_interfaces__::NamedTuple,
879+
__n_interfaces__::GraphPPL.StaticInt{$(length(ms_args))}
880+
)
881+
__interfaces__ = GraphPPL.prepare_interfaces(__model__, $ms_name, __lhs_interface__, __rhs_interfaces__)
882+
__context__ = GraphPPL.Context(__parent_context__, $ms_name)
883+
GraphPPL.copy_markov_blanket_to_child_context(__context__, __interfaces__)
884+
GraphPPL.add_composite_factor_node!(__model__, __parent_context__, __context__, $ms_name)
885+
__returnval__ = GraphPPL.add_terminated_submodel!(
886+
__model__, __context__, __options__, $ms_name, __interfaces__, __n_interfaces__
887+
)
888+
GraphPPL.returnval!(__context__, __returnval__)
889+
return __context__, __lhs_interface__
890+
end
891+
784892
function GraphPPL.add_terminated_submodel!(
785893
__model__::GraphPPL.Model,
786894
__context__::GraphPPL.Context,

0 commit comments

Comments
 (0)