Typing: make mypy . green (410 -> 0)

Typing-only campaign: dynamic attributes declared on Node and
HeadlessTransformer as bare class-body annotations; ortools CpModel
calls migrated to the snake_case API the stubs declare (runtime aliases
on 9.15, each verified before renaming; solver-side CamelCase kept
where snake_case equivalents are properties); the rest is annotations,
Optional corrections, and typing.cast. One type: ignore in the whole
tree. make lint passes; full Modal suite green except the pre-existing
staircase drift failure.

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

Author: physicsrob

examples/calculator_advanced.py

@@ -294,7 +294,7 @@ def _make_compressor(embedding: Embedding):
     }
 
     def compress(digits: List[Node]):
-        values = [Linear(d, value_proj, name="digit_value") for d in digits]
+        values: List[Node] = [Linear(d, value_proj, name="digit_value") for d in digits]
         total = sum_nodes(values)  # plain-number add of up to 11 digits, 0..99
         bucket = bool_to_01(in_range(total, add_const(total, 1.0), n_buckets))
         sum_digit = onehot_lookup(bucket, sum_table, zero_digit)

examples/calculator_scratchpad.py

@@ -70,7 +70,7 @@
 wrong-but-valid number.
 """
 
-from typing import Callable, Dict, List, Tuple
+from typing import Callable, Dict, List, Optional, Tuple, Union, cast
 
 import torch
 
@@ -580,7 +580,7 @@ def _emit_gathered_answer(
     col_onehot: Node,
     N: int,
     *,
-    sign_at: Callable[[int], Node] = None,
+    sign_at: Optional[Callable[[int], Node]] = None,
 ) -> List[Node]:
     """The trimmed, MSB-first answer, one slot per ``layout.n_answer``.
 
@@ -624,7 +624,7 @@ def _emit_gathered_answer(
             match_gain=_MATCH_GAIN,
         )
         if t == 0 and sign_at is not None:
-            answer.append(select(ge, picked, minus))
+            answer.append(select(cast(Node, ge), picked, minus))
         else:
             answer.append(select(compare(idx, thresh=N - 0.5), eos, picked))
     return answer
@@ -767,6 +767,7 @@ def _sub_op(
     # constant-depth; here the recurrence rides the decode axis, where
     # serial is the point. ---
     combine_table: Dict[torch.Tensor, torch.Tensor] = {}
+    nxt: Union[int, Node]
     for v in range(_CMP_W):
         for a in range(10):
             for b in range(10):

scripts/arithmetic_scaling.py

@@ -153,7 +153,7 @@ def critical_path_depth(outputs) -> int:
     Computed with an explicit stack (iterative post-order) so deep graphs do
     not overflow Python's recursion limit.
     """
-    memo = {}
+    memo: dict = {}
     for root in outputs:
         stack = [(root, False)]
         while stack:

scripts/calculator_layer_table.py

@@ -135,7 +135,7 @@ def collect_cell(impl_name: str, n: int) -> dict:
     try:
         out, embedding = impl.create_network_parts(max_digits=n)
     except Exception as exc:  # structural caps (slow planes, fact table)
-        row = {"n": n, "build_error": str(exc).splitlines()[0]}
+        row: dict = {"n": n, "build_error": str(exc).splitlines()[0]}
         if hasattr(impl, "n_params"):
             row["params_extrapolated"] = impl.n_params(n)
         if hasattr(impl, "compiled_layers"):

scripts/calculator_stats.py

@@ -19,6 +19,7 @@
 import json
 import os
 import tempfile
+from typing import cast
 
 import torch
 
@@ -63,7 +64,7 @@ def whole_calculator_stats() -> None:
             d_head=D_HEAD,
             verbose=False,
         )
-        sidecar = json.load(open(artifact.debug_path))
+        sidecar = json.load(open(cast(str, artifact.debug_path)))
 
     peak = _peak_residual_occupancy(sidecar)
     print(f"  layers              : {artifact.n_layers}")

scripts/diagnose_calculator_layer_gap.py

@@ -24,7 +24,7 @@
 import argparse
 import importlib
 from collections import Counter
-from typing import Dict, List, Tuple
+from typing import Any, Callable, Dict, List, Tuple, cast
 
 from torchwright.compiler.lower import lower
 from torchwright.compiler.forward.cpsat_scheduler import (
@@ -57,13 +57,13 @@ def longest_chain(gm, es: Dict[int, int]) -> List:
         if u.node_id in preds and v.node_id in preds:
             preds[v.node_id].append(u.node_id)
 
-    cur = max(es, key=es.get)
+    cur = max(es, key=cast(Callable[[int], int], es.get))
     chain = [cur]
     while True:
         candidates = [u for u in preds[cur] if es[u] >= es[cur] - 1]
         if not candidates:
             break
-        cur = max(candidates, key=es.get)
+        cur = max(candidates, key=cast(Callable[[int], int], es.get))
         chain.append(cur)
     return [by_id[i] for i in reversed(chain)]
 
@@ -196,7 +196,7 @@ def main() -> None:
     print(f"[{args.impl} n={args.n}] nonlinear-op depth: {depth}")
 
     if args.forensic:
-        lowered = lower(
+        lowered: Any = lower(
             out,
             collapse_univariate=True,
             collapse_pl=False,

scripts/ffn_equivalence.py

@@ -29,7 +29,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, cast
 
 import torch
 
@@ -151,7 +151,7 @@ def schedule_metrics(
         if output_node in computed:
             break
         attn_ops, mlp_ops, _biased = sched.schedule_layer(rmap, computed)
-        heads = sum(_count_heads_by_type(attn_ops, d_head).values())
+        heads = sum(_count_heads_by_type(cast(Any, attn_ops), d_head).values())
         per_layer_heads.append(heads)
         slots = sum(len(op.mlp_slots) for op in mlp_ops if op.mlp_slots)
         peak_hidden = max(peak_hidden, slots)
@@ -207,10 +207,10 @@ def schedule_trace(
             if op.op_type == "compute_ffn":
                 composites.append(
                     (
-                        op.node.annotation,
-                        op.node.d_output,
+                        cast(Node, op.node).annotation,
+                        cast(Node, op.node).d_output,
                         len(op.mlp_slots),
-                        op.node.node_id,
+                        cast(Node, op.node).node_id,
                     )
                 )
         trace.append({"layer": i, "hidden": hidden, "composites": composites})

scripts/investigate_log_precision.py

@@ -14,7 +14,7 @@
 
 import torch
 
-from torchwright.ops.relu.arithmetic_ops import log
+from torchwright.ops.relu.arithmetic_ops import log  # type: ignore[attr-defined]
 from torchwright.ops.inout_nodes import create_input
 
 

scripts/investigate_onehot_leak.py

@@ -38,6 +38,7 @@
 
 import torchwright.ops.swiglu.map_select as _map_select
 import torchwright.ops.swiglu.onehot_table as _onehot_table
+from torchwright.graph import Node
 from torchwright.graph.node import suppress_checks
 from torchwright.ops._math import _lookup_numeric_slack
 from torchwright.ops.inout_nodes import create_input
@@ -76,7 +77,7 @@ def __exit__(self, *exc):
     return _Ctx()
 
 
-def build_carry_chain(scale_override: Optional[float]) -> Dict[str, object]:
+def build_carry_chain(scale_override: Optional[float]) -> Dict[str, Node]:
     """total -> in_range(total, total+1, 61) -> bool_to_01 -> lookups.
 
     The exact shape of calculator_simple.multiply_digit_seqs' carry sweep,
@@ -110,7 +111,7 @@ def build_carry_chain(scale_override: Optional[float]) -> Dict[str, object]:
     }
 
 
-def build_times_table(scale_override: Optional[float]) -> Dict[str, object]:
+def build_times_table(scale_override: Optional[float]) -> Dict[str, Node]:
     """Two machine-built one-hot digits -> two-block times-table lookup.
 
     The multi-block (swiglu-lane) path of onehot_lookup, fed by the same
@@ -141,7 +142,7 @@ def build_times_table(scale_override: Optional[float]) -> Dict[str, object]:
 
 
 def sweep_carry(
-    nodes: Dict[str, object], deltas: List[float], device: torch.device
+    nodes: Dict[str, Node], deltas: List[float], device: torch.device
 ) -> Dict[str, torch.Tensor]:
     ts, dvals = [], []
     for t in range(N_SLOTS):
@@ -180,7 +181,7 @@ def sweep_carry(
 
 
 def sweep_times_table(
-    nodes: Dict[str, object], deltas: List[float], device: torch.device
+    nodes: Dict[str, Node], deltas: List[float], device: torch.device
 ) -> Dict[str, torch.Tensor]:
     rows: List[Tuple[int, int, float]] = []
     for da in range(10):

scripts/measure_fusion_opportunities.py

@@ -46,7 +46,7 @@
 import importlib
 import json
 from collections import Counter
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, cast
 
 from torchwright.compiler.collapse import scalar_sources
 from torchwright.compiler.graph_clone import topological_order
@@ -170,7 +170,7 @@ def _subgraph_details(
         if s is not None and s is not n:
             by_src.setdefault(s, []).append(n)
 
-    reports = []
+    reports: List[Dict[str, Any]] = []
     for s, members in by_src.items():
         mset = set(members)
         # Longest chain inside the subgraph, overall and critical-only.
@@ -271,7 +271,9 @@ def analyze(
     members = [n for n in order if src[n] is not None and src[n] is not n]
     lv2 = _collapsed_levels(order, src)
     depth2 = lv2[out]
-    subgraph_sizes = Counter(src[n].name or type(src[n]).__name__ for n in members)
+    subgraph_sizes = Counter(
+        cast(Node, src[n]).name or type(src[n]).__name__ for n in members
+    )
 
     # Feasibility-filtered variant: a collapsed univariate subgraph is one
     # FFN whose breakpoint grid must resolve the composed function over
@@ -304,7 +306,7 @@ def _feasible(s: Node) -> bool:
     ]
 
     # Families 1+2 combined.
-    lv12 = {}
+    lv12: Dict[Node, int] = {}
     for n in order:
         s = src[n]
         if s is not None and s is not n: