audit: align mixed-traffic compute with cost ledger

1 files changed, 15 insertions, 12 deletions

diff --git a/sdil/conv.py b/sdil/conv.py
index 0beb249..cc4c7c6 100644
--- a/sdil/conv.py
+++ b/sdil/conv.py
@@ -649,14 +649,15 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
     def mixed_elementwise_ops_per_example(self, rule=None):
         """Transparent arithmetic count beyond convolution/linear MACs.
 
-        Five operations per unit form traffic, affine prediction, raw, and
-        innovation. Norm matching additionally charges squares/reductions,
-        rescaling, and scalar norm arithmetic conservatively as five per unit.
+        Six operations per unit form traffic, affine prediction, raw, and
+        innovation in the executed tensor expression. Norm matching
+        additionally charges squares/reductions, rescaling, and scalar norm
+        arithmetic conservatively as five per unit.
         """
         rule = self.traffic_rule if rule is None else rule
         if rule not in ("raw", "matched", "innovation"):
             raise ValueError(f"unknown mixed-traffic rule: {rule}")
-        multiplier = 10 if rule == "matched" else 5
+        multiplier = 11 if rule == "matched" else 6
         return multiplier * self.mixed_units_per_example
 
     @property
@@ -704,7 +705,8 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
         }
 
     @torch.no_grad()
-    def mixed_apical_components(self, instruction, hiddens, rule):
+    def mixed_apical_components(self, instruction, hiddens, rule,
+                                compute_matched=False):
         """Return used, raw, innovation, matched, and traffic fields."""
         if rule not in ("raw", "matched", "innovation"):
             raise ValueError(f"unknown mixed-traffic rule: {rule}")
@@ -713,19 +715,20 @@ class CIFARKPMixedTrafficResNet(CIFARKPResNet):
         traffic = self.traffic_fields(hiddens)
         raw = []
         innovation = []
-        matched = []
+        matched = [] if (rule == "matched" or compute_matched) else None
         for signal, hidden, ordinary, slope, bias in zip(
                 instruction, hiddens, traffic,
                 self.P_traffic, self.P_traffic_bias):
             apical = signal + ordinary
             residual = apical - (slope * hidden + bias)
-            raw_norm = apical.flatten(1).norm(dim=1).clamp_min(1e-30)
-            residual_norm = residual.flatten(1).norm(dim=1)
-            scale = (residual_norm / raw_norm).reshape(
-                residual.shape[0], *([1] * (residual.ndim - 1)))
             raw.append(apical)
             innovation.append(residual)
-            matched.append(scale * apical)
+            if matched is not None:
+                raw_norm = apical.flatten(1).norm(dim=1).clamp_min(1e-30)
+                residual_norm = residual.flatten(1).norm(dim=1)
+                scale = (residual_norm / raw_norm).reshape(
+                    residual.shape[0], *([1] * (residual.ndim - 1)))
+                matched.append(scale * apical)
         choices = {"raw": raw, "matched": matched, "innovation": innovation}
         return {
             "used": choices[rule], "raw": raw, "innovation": innovation,
@@ -1335,7 +1338,7 @@ def conv_kp_mixed_traffic_alignment_report(net, x, y, rule):
                         - F.one_hot(y, net.n_classes).to(forward["logits"].dtype))
         instruction = net.hierarchical_teaching(output_error, forward)
         components = net.mixed_apical_components(
-            instruction, forward["hiddens"], rule)
+            instruction, forward["hiddens"], rule, compute_matched=True)
 
         def align(values):
             return [float(F.cosine_similarity(