Fix apply_op for tuple containers and keyword container inputs

[Max Freedom Pollard (MaxFreedomPollard)]

Summary

apply_op (exported from opto.trace) broadcasts an operator over a container of Nodes. Two of its supported container shapes raise instead of working. Both are one-line-class bugs with a two-line fix each, and both are covered by new regression tests.

1. Tuple containers raise TypeError

The list/tuple branch assigns into output element by element, which a tuple cannot do:

from opto.trace import node, apply_op
import opto.trace.operators as ops

apply_op(ops.add, (node("x1"), node("x2")),
                  (node("a1"), node("a2")),
                  (node("b1"), node("b2")))

TypeError: 'tuple' object does not support item assignment

The intent was already in the code — the branch ends with

if isinstance(output, tuple):
    output = tuple(output)

which is a no-op as written and is never reached for a tuple anyway. The fix records whether the output was a tuple, works on a list, and converts back at the end. Lists keep their current in-place behaviour (which test_apply_op.py relies on).

2. Keyword container inputs hit an assertion

The NodeContainer branch builds _kwargs with

kk: vv if isinstance(v, Node) else getattr(vv, k)

v is the output's attribute; vv is the keyword input. When an attribute of the output is a Node, every keyword input is forwarded whole instead of being indexed by attribute name, and the recursive call trips the admissible-type assertion:

from opto.trace import node, bundle, apply_op
from opto.trace.containers import NodeContainer

class C(NodeContainer):
    def __init__(self, x):
        self.x = node(x)

@bundle()
def concat(foo, bar):
    return foo + bar

apply_op(concat, C("seed"), foo=C("foo"), bar=C("bar"))

AssertionError

The positional path one line above already uses the input (x), and the list and dict branches already test isinstance(vv, Node). This makes the keyword path consistent with them; the same call passed positionally works today and returns "foobar".

Tests

tests/unit_tests/test_apply_op.py gains coverage for:

• a tuple container, standalone and nested inside a NodeContainer (result type, values, and that the inputs are wired in as parents)
• a list container still being updated in place
• keyword inputs whose values are containers, and keyword inputs mixing a container with a bare Node

Each new assertion fails on main and passes with the fix.

Verification

• python tests/unit_tests/run.py — 22/22 pass, exit code 0
• black==23.3.0 (the pinned pre-commit version) reports both files unchanged
• No API, signature, or dependency changes

[Max Freedom Pollard (MaxFreedomPollard)]

@microsoft-github-policy-service agree