Reformat with black 26.3.1

Mechanical: the unpinned dev-group black moved to 26.x, which formats
differently than the committed baseline; make lint's black step is
green again.

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

Author: physicsrob

scripts/dump_measure_noise.py

@@ -5,22 +5,27 @@
 
     make modal-run MODULE=scripts.dump_measure_noise > docs/op_noise_data.json
 """
+
 import sys
 from pathlib import Path
 
 from scripts.measure_op_noise import _measure_all, render_json
-from torchwright.compiler.forward.compile import forward_compile  # noqa: F401 pre-import
+from torchwright.compiler.forward.compile import (
+    forward_compile,
+)  # noqa: F401 pre-import
 
 
 def main():
     measurements = _measure_all()
     # Get commit SHA from env or fallback
     import subprocess
+
     try:
         sha = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode().strip()
     except Exception:
         sha = "unknown"
     import datetime
+
     now = datetime.datetime.utcnow().isoformat() + "Z"
     json_text = render_json(measurements, commit=sha, measured_at=now)
     # Print a separator so we can grep out the real JSON block

scripts/validate_multiply_2d_build_cost.py

@@ -44,9 +44,7 @@ def main() -> None:
     b = create_input("b", 1)
     tracemalloc.start()
     t0 = time.perf_counter()
-    node = multiply_2d(
-        a, b, max_abs1=max_abs, max_abs2=max_abs, step1=step, step2=step
-    )
+    node = multiply_2d(a, b, max_abs1=max_abs, max_abs2=max_abs, step1=step, step2=step)
     elapsed = time.perf_counter() - t0
     _cur, peak = tracemalloc.get_traced_memory()
     tracemalloc.stop()

tests/compile/forward/test_cpsat_knobs.py

@@ -650,9 +650,7 @@ def schedule_layer(self, *args, **kwargs):
     assert list(tmp_path.glob("*.json")) == []
 
 
-def test_schedule_cache_stores_once_after_replay_and_not_on_hit(
-    tmp_path, monkeypatch
-):
+def test_schedule_cache_stores_once_after_replay_and_not_on_hit(tmp_path, monkeypatch):
     """The positive mirror: a successful non-cached solve stores exactly once
     (after replay), and a cache hit does not store again."""
     from torchwright.compiler.forward import compile as compile_mod

tests/compile/forward/test_cpsat_pin_cancels.py

@@ -39,7 +39,6 @@
 from torchwright.compiler.forward.scheduling_policy import SchedulingPolicy
 from torchwright.compiler.lower import lower
 
-
 # ---------------------------------------------------------------------------
 # The six example graphs (matches test_cpsat_sym1.py / test_cpsat_snapshot.py).
 # ---------------------------------------------------------------------------
@@ -60,7 +59,13 @@ def _build_bucketed_argmin():
     value = InputNode("baib_value", vw, value_range=(-100.0, 100.0))
     return attend_argmin_above_in_bucket(
         create_rope_config(d_head=d_head, max_positions=512),
-        score, validity, kb, above, qb, th, value,
+        score,
+        validity,
+        kb,
+        above,
+        qb,
+        th,
+        value,
     )
 
 
@@ -77,12 +82,22 @@ def _example_specs():
         # calculator: d pinned at the pre-unification test width (the canonical
         # publish D_MODEL=8192 would make this schedule test enormous for
         # nothing); d_head follows the module — it is baked into the graph.
-        "calculator": (lambda: calculator_simple.create_network_parts()[0], 1024, calculator_simple.D_HEAD),
+        "calculator": (
+            lambda: calculator_simple.create_network_parts()[0],
+            1024,
+            calculator_simple.D_HEAD,
+        ),
         "caesar": (lambda: caesar_cipher.create_network_parts()[0], 512, 16),
-        "sort_digits": (lambda: sort_digits_v1.create_network_parts()[0], sort_digits_v1.D_MODEL, sort_digits_v1.D_HEAD),
+        "sort_digits": (
+            lambda: sort_digits_v1.create_network_parts()[0],
+            sort_digits_v1.D_MODEL,
+            sort_digits_v1.D_HEAD,
+        ),
         "fibonacci": (lambda: fibonacci.create_network_parts()[0], 512, 16),
         "binary_increment": (
-            lambda: binary_increment.create_network_parts()[0], 256, 16
+            lambda: binary_increment.create_network_parts()[0],
+            256,
+            16,
         ),
         "bucketed_argmin": (_build_bucketed_argmin, 512, 32),
     }
@@ -129,9 +144,9 @@ def test_default_build_is_pinned(name):
     on = _proto_text(build_cpsat_model(node, _pin_cancels=True, **cfg))
     assert default == on, f"{name}: default proto differs from pinned build"
     assert "parked" not in default, f"{name}: default build has parked vars"
-    assert "pin_attn" in default or "pin_mlp" in default, (
-        f"{name}: default build posts no pin aux vars (pin is dead)"
-    )
+    assert (
+        "pin_attn" in default or "pin_mlp" in default
+    ), f"{name}: default build posts no pin aux vars (pin is dead)"
 
 
 @pytest.mark.parametrize("name", _NAMES)
@@ -147,9 +162,9 @@ def test_knob_off_reproduces_legacy_model(name):
     off = _proto_text(build_cpsat_model(node, _pin_cancels=False, **cfg))
     assert off != default, f"{name}: knob-off proto identical to default"
     assert "parked" in off, f"{name}: legacy build has no parked vars"
-    assert "pin_attn" not in off and "pin_mlp" not in off, (
-        f"{name}: legacy build still posts pin aux vars"
-    )
+    assert (
+        "pin_attn" not in off and "pin_mlp" not in off
+    ), f"{name}: legacy build still posts pin aux vars"
 
 
 # ---------------------------------------------------------------------------
@@ -168,8 +183,13 @@ def test_knob_off_reproduces_legacy_model(name):
 
 def _solve_cfg(d, d_head):
     return dict(
-        d=d, d_head=d_head, d_hidden=d, max_layers=100,
-        time_budget_s=60.0, policy=SchedulingPolicy(), tighten_domains=True,
+        d=d,
+        d_head=d_head,
+        d_hidden=d,
+        max_layers=100,
+        time_budget_s=60.0,
+        policy=SchedulingPolicy(),
+        tighten_domains=True,
     )
 
 
@@ -190,8 +210,13 @@ def test_pinned_solution_valid_in_unpinned_model(name, d):
     )
 
     built = build_cpsat_model(
-        low, d=d, d_head=d_head, d_hidden=d, max_layers=100,
-        policy=SchedulingPolicy(), tighten_domains=True,
+        low,
+        d=d,
+        d_head=d_head,
+        d_hidden=d,
+        max_layers=100,
+        policy=SchedulingPolicy(),
+        tighten_domains=True,
         _pin_cancels=False,  # the legacy model is the verification target
     )
     for nid, L in asg.node_to_layer.items():
@@ -228,12 +253,12 @@ def _solve(pin):
 
     off_asg, off_stats = _solve(False)
     on_asg, on_stats = _solve(True)
-    assert off_asg is not None and off_stats.is_optimal, (
-        f"{name} d={d}: unpinned not optimal in budget ({off_stats.status_name})"
-    )
-    assert on_asg is not None and on_stats.is_optimal, (
-        f"{name} d={d}: pinned not optimal in budget ({on_stats.status_name})"
-    )
+    assert (
+        off_asg is not None and off_stats.is_optimal
+    ), f"{name} d={d}: unpinned not optimal in budget ({off_stats.status_name})"
+    assert (
+        on_asg is not None and on_stats.is_optimal
+    ), f"{name} d={d}: pinned not optimal in budget ({on_stats.status_name})"
     assert on_asg.n_layers >= off_asg.n_layers, (
         f"{name} d={d}: pinned optimum {on_asg.n_layers} beats unpinned "
         f"{off_asg.n_layers} — impossible for a pure restriction"
@@ -267,9 +292,9 @@ def test_pin_reaches_snapshot_path():
     )
     live_default = _proto_text(build_cpsat_model(node, **cfg))
     assert snap_default != snap_off, "knob dead on the snapshot path"
-    assert snap_default == live_default, (
-        "default (pinned) snapshot proto differs from default live"
-    )
+    assert (
+        snap_default == live_default
+    ), "default (pinned) snapshot proto differs from default live"
 
 
 # ---------------------------------------------------------------------------
@@ -298,9 +323,9 @@ def test_full_hint_with_pin_passes_strict_validation():
         _pin_cancels=True,
         **cfg,
     )
-    assert asg is not None, (
-        f"pinned solve with full hint found nothing ({stats.status_name})"
-    )
+    assert (
+        asg is not None
+    ), f"pinned solve with full hint found nothing ({stats.status_name})"
 
 
 # ---------------------------------------------------------------------------
@@ -315,8 +340,12 @@ def test_default_pinned_compile_replays_clean():
     build, d, d_head = _example_specs()["caesar"]
     torch.manual_seed(0)
     net = forward_compile(
-        d=d, d_head=d_head, output_node=build(), device="cpu",
-        verbose=False, optimize=1,
+        d=d,
+        d_head=d_head,
+        output_node=build(),
+        device="cpu",
+        verbose=False,
+        optimize=1,
     )
     # A raise above would be the replay-depth tripwire (or any compile error)
     # firing; a clean return with real layers is the tripwire staying silent.
@@ -327,7 +356,13 @@ def test_knob_off_compile_replays_clean():
     build, d, d_head = _example_specs()["caesar"]
     torch.manual_seed(0)
     net = forward_compile(
-        d=d, d_head=d_head, output_node=build(), device="cpu",
-        verbose=False, optimize=1, _pin_cancels=False, _force_resolve=True,
+        d=d,
+        d_head=d_head,
+        output_node=build(),
+        device="cpu",
+        verbose=False,
+        optimize=1,
+        _pin_cancels=False,
+        _force_resolve=True,
     )
     assert len(net.layers) > 0

tests/compile/forward/test_cpsat_snapshot.py

@@ -32,7 +32,6 @@
 )
 from torchwright.compiler.forward.scheduling_policy import SchedulingPolicy
 
-
 # ---------------------------------------------------------------------------
 # The six example graphs (matches scripts/intralayer_example_sweep.py)
 # ---------------------------------------------------------------------------
@@ -53,7 +52,13 @@ def _build_bucketed_argmin():
     value = InputNode("baib_value", vw, value_range=(-100.0, 100.0))
     return attend_argmin_above_in_bucket(
         create_rope_config(d_head=d_head, max_positions=512),
-        score, validity, kb, above, qb, th, value,
+        score,
+        validity,
+        kb,
+        above,
+        qb,
+        th,
+        value,
     )
 
 
@@ -70,12 +75,22 @@ def _example_specs():
         # calculator: d pinned at the pre-unification test width (the canonical
         # publish D_MODEL=8192 would make this schedule test enormous for
         # nothing); d_head follows the module — it is baked into the graph.
-        "calculator": (lambda: calculator_simple.create_network_parts()[0], 1024, calculator_simple.D_HEAD),
+        "calculator": (
+            lambda: calculator_simple.create_network_parts()[0],
+            1024,
+            calculator_simple.D_HEAD,
+        ),
         "caesar": (lambda: caesar_cipher.create_network_parts()[0], 512, 16),
-        "sort_digits": (lambda: sort_digits_v1.create_network_parts()[0], sort_digits_v1.D_MODEL, sort_digits_v1.D_HEAD),
+        "sort_digits": (
+            lambda: sort_digits_v1.create_network_parts()[0],
+            sort_digits_v1.D_MODEL,
+            sort_digits_v1.D_HEAD,
+        ),
         "fibonacci": (lambda: fibonacci.create_network_parts()[0], 512, 16),
         "binary_increment": (
-            lambda: binary_increment.create_network_parts()[0], 256, 16
+            lambda: binary_increment.create_network_parts()[0],
+            256,
+            16,
         ),
         "bucketed_argmin": (_build_bucketed_argmin, 512, 32),
     }
@@ -128,8 +143,12 @@ def test_canonicalized_snapshot_rebuilds_a_valid_model():
     node, d, d_head = _build("caesar")
     cfg = dict(d=d, d_head=d_head, d_hidden=d, max_layers=40)
     problem = snapshot_from_graph_model(build_graph_model(node)).canonicalized(node)
-    a = _proto_text(build_model_from_snapshot(SchedulingProblem.loads(problem.dumps()), **cfg))
-    b = _proto_text(build_model_from_snapshot(SchedulingProblem.loads(problem.dumps()), **cfg))
+    a = _proto_text(
+        build_model_from_snapshot(SchedulingProblem.loads(problem.dumps()), **cfg)
+    )
+    b = _proto_text(
+        build_model_from_snapshot(SchedulingProblem.loads(problem.dumps()), **cfg)
+    )
     assert a == b  # canonical rebuild is deterministic across loads
 
 
@@ -144,7 +163,11 @@ def test_hinted_build_proto_identical():
     # would make this proto comparison vacuous on the widening path.
     node, d, d_head = _build("caesar")
     cfg = dict(
-        d=d, d_head=d_head, d_hidden=d, max_layers=40, tighten_domains=True,
+        d=d,
+        d_head=d_head,
+        d_hidden=d,
+        max_layers=40,
+        tighten_domains=True,
         _pin_cancels=False,
     )
     # A feasible schedule -> hint dicts that exercise the cancel-window widening.
@@ -158,9 +181,7 @@ def test_hinted_build_proto_identical():
     problem = SchedulingProblem.loads(
         snapshot_from_graph_model(build_graph_model(node)).dumps()
     )
-    snap = _proto_text(
-        build_model_from_snapshot(problem, diagnostic_hint=hint, **cfg)
-    )
+    snap = _proto_text(build_model_from_snapshot(problem, diagnostic_hint=hint, **cfg))
     assert snap == live
 
 
@@ -186,8 +207,13 @@ def test_frozen_hint_roundtrips():
 def test_solve_from_snapshot_matches_live_solve():
     node, d, d_head = _build("caesar")
     kw = dict(
-        d=d, d_head=d_head, d_hidden=d, max_layers=40, time_budget_s=30.0,
-        tighten_domains=True, policy=SchedulingPolicy(),
+        d=d,
+        d_head=d_head,
+        d_hidden=d,
+        max_layers=40,
+        time_budget_s=30.0,
+        tighten_domains=True,
+        policy=SchedulingPolicy(),
         solver_params={"random_seed": 1},
     )
     live_asg, live_stats = solve_schedule(node, **kw)
@@ -208,8 +234,13 @@ def test_solve_from_snapshot_matches_live_solve():
 def test_save_load_with_fingerprint_gate(tmp_path):
     node, d, d_head = _build("binary_increment")
     geom = dict(
-        d=d, d_head=d_head, d_hidden=d, flex_routing=True, cancel_slack=2,
-        policy=SchedulingPolicy(), max_layers=60,
+        d=d,
+        d_head=d_head,
+        d_hidden=d,
+        flex_routing=True,
+        cancel_slack=2,
+        policy=SchedulingPolicy(),
+        max_layers=60,
     )
     cp = 5  # arbitrary stand-in for the lowered critical path in this unit test
     problem = (