compile_headless: one signature, one forward_compile call, schedule parity

New convention: compile_headless(graph, pos_encoding, *, ...) where
graph is an output Node (node mode: natural-column gather) or an io
dict (overlay mode: outputs land at input columns). The
isinstance(first_arg, PosEncoding) sniffing, the io=/output_node=
keywords, and the forked ~80-line _compile_headless_legacy are gone;
the two modes share one forward_compile call (node mode passes
overlays=None — identical to the old call) with thin pre/post
branches. A PosEncoding first arg raises a dedicated TypeError naming
the new order (PosEncoding IS a Node, so this check precedes the Node
branch).

Schedule parity: optimize and assume_zero_init now thread to
forward_compile (defaults 0/False == forward_compile's own defaults,
so existing behavior is unchanged) — the in-process debug backend can
finally reproduce a production optimize=2 schedule. Also max_layers
100->400 (pure cap at optimize=0) and verbose True->False, matching
the exporters.

Call sites flipped to (io, pos) order: test_io_api.py (plus the three
API-shape tests rewritten for the new convention),
test_residual_nonzero_init.py, test_probe.py. Legacy (node, pos)
sites — including probe_graph and doom's compiled_model — work
unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: physicsrob

tests/compile/forward/test_io_api.py

@@ -1,6 +1,6 @@
-"""Tests for the new io API in compile_headless.
+"""Tests for the io-dict form of compile_headless.
 
-The io API enables overlaid I/O where output values land at input columns
+The io form enables overlaid I/O where output values land at input columns
 via delta transfer, enabling autoregressive feedback where the transformer
 output IS the next input.
 """
@@ -27,8 +27,8 @@ def test_io_identity():
     x = create_input(4)
 
     module = compile_headless(
+        {"x": (x, x)},  # output = input
         pos,
-        io={"x": (x, x)},  # output = input
         d=64,
         verbose=False,
     )
@@ -45,8 +45,8 @@ def test_io_overlaid_single():
     y = multiply_const(x, 2.0)  # y = 2*x
 
     module = compile_headless(
+        {"x": (x, y)},  # x -> y, overlaid
         pos,
-        io={"x": (x, y)},  # x -> y, overlaid
         d=64,
         verbose=False,
     )
@@ -65,8 +65,8 @@ def test_io_overlaid_add_const():
     y = add_const(x, 10.0)  # y = x + 10
 
     module = compile_headless(
+        {"x": (x, y)},
         pos,
-        io={"x": (x, y)},
         d=64,
         verbose=False,
     )
@@ -88,11 +88,11 @@ def test_io_overlaid_multiple():
     out_b = add_const(b, 1.0)
 
     module = compile_headless(
-        pos,
-        io={
+        {
             "a": (a, out_a),  # Overlaid
             "b": (b, out_b),  # Overlaid
         },
+        pos,
         d=64,
         verbose=False,
     )
@@ -115,11 +115,11 @@ def test_io_input_only_not_in_output():
     out_b = multiply_const(b, 2.0)
 
     module = compile_headless(
-        pos,
-        io={
+        {
             "a": (a, None),  # Input-only: NOT in output
             "b": (b, out_b),  # Overlaid
         },
+        pos,
         d=64,
         verbose=False,
     )
@@ -142,11 +142,11 @@ def test_io_with_overflow():
     out_extra = create_literal_value(torch.tensor([99.0]))
 
     module = compile_headless(
-        pos,
-        io={
+        {
             "a": (a, out_a),  # Overlaid at columns [0:2]
             "extra": (None, out_extra),  # Overflow at column [2]
         },
+        pos,
         d=64,
         verbose=False,
     )
@@ -166,8 +166,8 @@ def test_io_autoregressive():
     next_state = add_const(state, 1.0)
 
     module = compile_headless(
+        {"state": (state, next_state)},
         pos,
-        io={"state": (state, next_state)},
         d=64,
         verbose=False,
     )
@@ -191,11 +191,11 @@ def test_io_uses_input_value():
     out_xy = add(x, y)
 
     module = compile_headless(
-        pos,
-        io={
+        {
             "x": (x, out_xy),  # Overlaid: output x+y at x's columns
             "y": (y, None),  # Input-only
         },
+        pos,
         d=64,
         verbose=False,
     )
@@ -221,15 +221,15 @@ def test_io_width_mismatch_error():
     out_a = create_literal_value(torch.zeros(2))  # Wrong width
 
     with pytest.raises(ValueError, match="width"):
-        compile_headless(pos, io={"a": (a, out_a)}, d=64, verbose=False)
+        compile_headless({"a": (a, out_a)}, pos, d=64, verbose=False)
 
 
 def test_io_empty_tuple_error():
     """Each io entry must have at least one node."""
     pos = create_pos_encoding()
 
     with pytest.raises(ValueError, match="both input and output as None"):
-        compile_headless(pos, io={"a": (None, None)}, d=64, verbose=False)
+        compile_headless({"a": (None, None)}, pos, d=64, verbose=False)
 
 
 def test_io_no_output_error():
@@ -238,21 +238,21 @@ def test_io_no_output_error():
     x = create_input(4)
 
     with pytest.raises(ValueError, match="at least one output"):
-        compile_headless(pos, io={"x": (x, None)}, d=64, verbose=False)
+        compile_headless({"x": (x, None)}, pos, d=64, verbose=False)
 
 
 # ---------------------------------------------------------------------------
-# Legacy API compatibility tests
+# Calling-convention tests
 # ---------------------------------------------------------------------------
 
 
-def test_legacy_api_still_works():
-    """Legacy output_node API should still work."""
+def test_node_mode():
+    """A bare output Node compiles in node mode (natural-column gather)."""
     pos = create_pos_encoding()
-    x = create_input("x", 4)  # Legacy: named input
+    x = create_input("x", 4)  # node mode: named input
     y = multiply_const(x, 2.0)
 
-    module = compile_headless(pos, output_node=y, d=64, verbose=False)
+    module = compile_headless(y, pos, d=64, verbose=False)
 
     inp = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
     out = module(inp)
@@ -261,22 +261,23 @@ def test_legacy_api_still_works():
     assert torch.allclose(out, expected, atol=0.1)
 
 
-def test_cannot_specify_both_io_and_output_node():
-    """Cannot specify both io and output_node."""
+def test_output_node_keyword_rejected():
+    """The removed output_node= keyword raises TypeError."""
     pos = create_pos_encoding()
     x = create_input(4)
     y = multiply_const(x, 2.0)
 
-    with pytest.raises(ValueError, match="Cannot specify both"):
-        compile_headless(pos, io={"x": (x, y)}, output_node=y, d=64, verbose=False)
+    with pytest.raises(TypeError):
+        compile_headless({"x": (x, y)}, pos, output_node=y, d=64, verbose=False)
 
 
-def test_must_specify_io_or_output_node():
-    """Must specify either io or output_node."""
+def test_pos_encoding_first_rejected_with_hint():
+    """The old (pos_encoding, io) order raises a TypeError naming the fix."""
     pos = create_pos_encoding()
+    x = create_input(4)
 
-    with pytest.raises(ValueError, match="Either io or output_node"):
-        compile_headless(pos, d=64, verbose=False)
+    with pytest.raises(TypeError, match="graph, pos_encoding"):
+        compile_headless(pos, {"x": (x, x)}, d=64, verbose=False)
 
 
 # ---------------------------------------------------------------------------
@@ -292,8 +293,8 @@ def test_anonymous_input_with_io():
     assert x.name == ""  # Verify it's anonymous
 
     module = compile_headless(
+        {"my_input": (x, x)},  # Name comes from io key
         pos,
-        io={"my_input": (x, x)},  # Name comes from io key
         d=64,
         verbose=False,
     )

tests/compile/forward/test_residual_nonzero_init.py

@@ -73,7 +73,7 @@ def test_nonzero_init_intermediate_col(monkeypatch):
     x = create_input(4)
     y = multiply_const(x, 2.0)
 
-    module = compile_headless(pos, io={"x": (x, y)}, d=64, verbose=False)
+    module = compile_headless({"x": (x, y)}, pos, d=64, verbose=False)
     inp = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
     expected = torch.tensor([[2.0, 4.0, 6.0, 8.0]])
 
@@ -100,8 +100,8 @@ def test_nonzero_init_overflow_output(monkeypatch):
     y = add_const(x, 10.0)
 
     module = compile_headless(
+        {"a": (x, None), "b": (None, y)},
         pos,
-        io={"a": (x, None), "b": (None, y)},
         d=64,
         verbose=False,
     )
@@ -134,7 +134,7 @@ def test_nonzero_init_mlp_path(monkeypatch):
     x = create_input(4)
     y = relu_add(x, multiply_const(x, -1.0))  # ReLU(x) + ReLU(-x) = |x|
 
-    module = compile_headless(pos, io={"x": (x, y)}, d=128, verbose=False)
+    module = compile_headless({"x": (x, y)}, pos, d=128, verbose=False)
     inp = torch.tensor([[1.0, -2.0, 3.0, -4.0]])
     expected = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
 

tests/debug/test_probe.py

@@ -73,8 +73,8 @@ def test_probe_residual_reads_intermediate_node():
     y = multiply_const(x, 3.0)
     z = add_const(y, 1.0)
     module = compile_headless(
+        {"x": (x, z)},
         pos,
-        io={"x": (x, z)},
         d=64,
         verbose=False,
     )
@@ -115,8 +115,8 @@ def test_probe_attention_captures_softmax_weights():
         output_matrix=torch.randn(4, 4),
     )
     module = compile_headless(
+        {"x": (x, attn)},
         pos,
-        io={"x": (x, attn)},
         d=256,
         d_head=16,
         verbose=False,
@@ -161,8 +161,8 @@ def test_probe_layer_diff_drift_and_sentinel():
     x = create_input("x", 2)
     y = multiply_const(x, 3.0)
     module = compile_headless(
+        {"x": (x, y)},
         pos,
-        io={"x": (x, y)},
         d=64,
         verbose=False,
     )