From 423161204c1f736329effec68f8eb042adaa2c15 Mon Sep 17 00:00:00 2001 From: MaxFreedomPollard <272618364+MaxFreedomPollard@users.noreply.github.com> Date: Fri, 4 Sep 2026 10:38:00 -0400 Subject: [PATCH] Fix apply_op for tuple containers and keyword container inputs apply_op broadcasts an operator over containers of Nodes, but two cases in opto/trace/broadcast.py never worked: - A tuple container raised "TypeError: 'tuple' object does not support item assignment", because the loop assigned into the tuple in place. The trailing `if isinstance(output, tuple): output = tuple(output)` was therefore unreachable. Accumulate into a list and convert back. - Keyword inputs against a Node-valued attribute of a NodeContainer tested `isinstance(v, Node)` (the output's attribute) instead of `isinstance(vv, Node)` (the keyword input), so a container passed by keyword was forwarded whole instead of being indexed by attribute, tripping the admissible-type assertion. The positional path already did the right thing. Adds regression tests covering a tuple container standalone and nested in a NodeContainer, a list and a dict; a bare Node broadcast against a tuple; the existing in-place list behaviour; and keyword inputs whose values are containers, including one mixed with a bare Node. --- opto/trace/broadcast.py | 8 +++- tests/unit_tests/test_apply_op.py | 79 ++++++++++++++++++++++++++++++- 2 files changed, 84 insertions(+), 3 deletions(-) diff --git a/opto/trace/broadcast.py b/opto/trace/broadcast.py index f157aa1f..94348cdd 100644 --- a/opto/trace/broadcast.py +++ b/opto/trace/broadcast.py @@ -67,13 +67,17 @@ def admissible_type(x, base): assert all( isinstance(x, Node) or len(output) == len(x) for x in inputs ), f"output {output} and inputs {inputs} are of different lengths." + # Tuples are immutable, so we accumulate into a list and convert back. + is_tuple = isinstance(output, tuple) + if is_tuple: + output = list(output) for k in range(len(output)): _args = [x if isinstance(x, Node) else x[k] for x in args] _kwargs = { kk: vv if isinstance(vv, Node) else vv[k] for kk, vv in kwargs.items() } output[k] = apply_op(op, output[k], *_args, **_kwargs) - if isinstance(output, tuple): + if is_tuple: output = tuple(output) elif isinstance(output, dict): @@ -88,7 +92,7 @@ def admissible_type(x, base): for k, v in output.__dict__.items(): _args = [x if isinstance(x, Node) else getattr(x, k) for x in args] _kwargs = { - kk: vv if isinstance(v, Node) else getattr(vv, k) + kk: vv if isinstance(vv, Node) else getattr(vv, k) for kk, vv in kwargs.items() } new_v = apply_op(op, v, *_args, **_kwargs) diff --git a/tests/unit_tests/test_apply_op.py b/tests/unit_tests/test_apply_op.py index dc64fa3a..25ebfd59 100644 --- a/tests/unit_tests/test_apply_op.py +++ b/tests/unit_tests/test_apply_op.py @@ -1,4 +1,4 @@ -from opto.trace import node +from opto.trace import node, bundle from opto.trace.broadcast import apply_op from opto.trace.containers import NodeContainer import opto.trace.operators as ops @@ -53,3 +53,80 @@ def __init__(self, x, v): assert foobar.dict_x["x"][0]["bar"].data == "bar1" assert foobar.dict_x["x"][1]["foo"].data == "foo2" assert foobar.dict_x["x"][1]["bar"].data == "bar2" + + +# Test tuple containers +class TupleContainer(NodeContainer): + def __init__(self, x): + self.tuple_x = (node(x + "1"), node(x + "2")) + + +foo_tuple = (node("foo1"), node("foo2")) +bar_tuple = (node("bar1"), node("bar2")) +out_tuple = apply_op(ops.add, (node("seed1"), node("seed2")), foo_tuple, bar_tuple) +assert isinstance(out_tuple, tuple) +assert out_tuple[0].data == "foo1bar1" +assert out_tuple[1].data == "foo2bar2" +assert foo_tuple[0] in out_tuple[0].parents and bar_tuple[0] in out_tuple[0].parents + +# A tuple nested inside a NodeContainer is replaced by the new tuple +foo_tc, bar_tc = TupleContainer("foo"), TupleContainer("bar") +out_tc = apply_op(ops.add, TupleContainer("seed"), foo_tc, bar_tc) +assert isinstance(out_tc.tuple_x, tuple) +assert out_tc.tuple_x[0].data == "foo1bar1" +assert out_tc.tuple_x[1].data == "foo2bar2" + +# A tuple nested in a list or a dict is likewise replaced by the new tuple +out_nested = apply_op( + ops.add, + [(node("seed1"), node("seed2"))], + [(node("foo1"), node("foo2"))], + [(node("bar1"), node("bar2"))], +) +assert isinstance(out_nested[0], tuple) +assert out_nested[0][0].data == "foo1bar1" +assert out_nested[0][1].data == "foo2bar2" + +out_nested = apply_op( + ops.add, + {"k": (node("seed1"), node("seed2"))}, + {"k": (node("foo1"), node("foo2"))}, + {"k": (node("bar1"), node("bar2"))}, +) +assert isinstance(out_nested["k"], tuple) +assert out_nested["k"][0].data == "foo1bar1" +assert out_nested["k"][1].data == "foo2bar2" + +# A bare Node is broadcast against every element of a tuple +z = node("Z") +out_tuple = apply_op(ops.add, (node("seed1"), node("seed2")), foo_tuple, z) +assert isinstance(out_tuple, tuple) +assert out_tuple[0].data == "foo1Z" +assert out_tuple[1].data == "foo2Z" + +# A list container is still updated in place +foo_list = [node("foo1"), node("foo2")] +bar_list = [node("bar1"), node("bar2")] +seed_list = [node("seed1"), node("seed2")] +out_list = apply_op(ops.add, seed_list, foo_list, bar_list) +assert out_list is seed_list +assert out_list[0].data == "foo1bar1" +assert out_list[1].data == "foo2bar2" + + +# Test keyword inputs against a Node-valued attribute +@bundle() +def concat(foo, bar): + return foo + bar + + +foo_sub, bar_sub = SubContainer("foo"), SubContainer("bar") +out_sub = apply_op(concat, SubContainer("seed"), foo=foo_sub, bar=bar_sub) +assert out_sub.y.data == "foobar" +assert foo_sub.y in out_sub.y.parents and bar_sub.y in out_sub.y.parents + +# Keyword inputs may mix bare Nodes with containers +bare = node("Z") +out_sub = apply_op(concat, SubContainer("seed"), foo=SubContainer("foo"), bar=bare) +assert out_sub.y.data == "fooZ" +assert bare in out_sub.y.parents