1import json
  2import random
  3import time
  4from collections import Counter
  5from itertools import combinations
  6
  7import numpy as np
  8import torch
  9
 10SEED = 597
 11random.seed(SEED)
 12np.random.seed(SEED)
 13torch.manual_seed(SEED)
 14
 15
 16def choose_device():
 17    if torch.cuda.is_available():
 18        try:
 19            torch.zeros(1, device="cuda")
 20            return "cuda"
 21        except Exception:
 22            pass
 23    return "cpu"
 24
 25
 26class Primitive:
 27    def __init__(self, name, support, kind, weight, bias, scale):
 28        self.name = name
 29        self.support = frozenset(support)
 30        self.kind = kind
 31        self.weight = weight
 32        self.bias = bias
 33        self.scale = scale
 34
 35    def apply(self, x):
 36        # Every primitive reads/writes only its declared coordinates.
 37        idx = sorted(self.support)
 38        y = x.clone()
 39        z = x[..., idx]
 40        if self.kind == "affine":
 41            z = z @ self.weight + self.bias
 42        elif self.kind == "tanh_affine":
 43            z = torch.tanh(z @ self.weight + self.bias)
 44        elif self.kind == "scale":
 45            z = z * self.scale
 46        y[..., idx] = z
 47        return y
 48
 49
 50def build_library(dim=12, groups=4, per_group=2, device="cpu"):
 51    ops = {}
 52    # Three independently parameterized operators per group. Operators within
 53    # one group need not commute; operators from different groups do commute.
 54    for g in range(groups):
 55        support = list(range(g * per_group, (g + 1) * per_group))
 56        for j, kind in enumerate(("affine", "tanh_affine", "scale")):
 57            name = f"g{g}_{j}"
 58            if kind == "scale":
 59                weight = bias = None
 60                scale = torch.rand(per_group, device=device) * 0.8 + 0.6
 61            else:
 62                weight = torch.randn(per_group, per_group, device=device) * 0.35
 63                bias = torch.randn(per_group, device=device) * 0.1
 64                scale = None
 65            ops[name] = Primitive(name, support, kind, weight, bias, scale)
 66    return ops
 67
 68
 69def commute(a, b, ops):
 70    return bool(ops[a].support.isdisjoint(ops[b].support))
 71
 72
 73def canonical(word, order, ops):
 74    # Bubble-sort only inverted adjacent commuting letters. This is the
 75    # requested local rewrite system for a fixed total order.
 76    rank = {name: i for i, name in enumerate(order)}
 77    out = list(word)
 78    changed = True
 79    while changed:
 80        changed = False
 81        for i in range(len(out) - 1):
 82            a, b = out[i], out[i + 1]
 83            if rank[a] > rank[b] and commute(a, b, ops):
 84                out[i], out[i + 1] = b, a
 85                changed = True
 86    return tuple(out)
 87
 88
 89def run_word(word, ops, x):
 90    y = x
 91    for name in word:
 92        y = ops[name].apply(y)
 93    return y
 94
 95
 96def main():
 97    device = choose_device()
 98    ops = build_library(device=device)
 99    order = sorted(ops)
100    names = list(ops)
101    x = torch.randn(32, 12, device=device)
102
103    # Core mathematical sanity check: every proposed adjacent commuting swap
104    # should preserve the complete tensor output to numerical precision.
105    swap_errors = []
106    noncommuting_differences = []
107    rng = random.Random(SEED + 1)
108    for _ in range(300):
109        a, b = rng.sample(names, 2)
110        prefix = [rng.choice(names) for _ in range(3)]
111        suffix = [rng.choice(names) for _ in range(3)]
112        w1 = prefix + [a, b] + suffix
113        w2 = prefix + [b, a] + suffix
114        err = (run_word(w1, ops, x) - run_word(w2, ops, x)).abs().max().item()
115        if commute(a, b, ops):
116            swap_errors.append(err)
117        else:
118            noncommuting_differences.append(err)
119    max_commuting_error = max(swap_errors)
120    median_noncommuting_difference = float(np.median(noncommuting_differences))
121
122    # Canonical words are functionally equal to their original words.
123    canonical_errors = []
124    for _ in range(100):
125        word = tuple(rng.choice(names) for _ in range(12))
126        cw = canonical(word, order, ops)
127        canonical_errors.append((run_word(word, ops, x) - run_word(cw, ops, x)).abs().max().item())
128
129    # Search/cache toy: fixed 500 proposed traces. Baseline evaluates all
130    # strings; canonicalized search evaluates one forward per equivalence key.
131    proposals = []
132    for _ in range(50):
133        base = [rng.choice(names) for _ in range(12)]
134        for _ in range(10):
135            v = list(base)
136            for _ in range(20):
137                i = rng.randrange(len(v) - 1)
138                if commute(v[i], v[i + 1], ops):
139                    v[i], v[i + 1] = v[i + 1], v[i]
140            proposals.append(tuple(v))
141    baseline_unique = len(set(proposals))
142    canonical_keys = [canonical(w, order, ops) for w in proposals]
143    idea_unique = len(set(canonical_keys))
144    # Use a deterministic target and an actual scalar objective so cache hits
145    # represent the same architecture/function, not merely string equality.
146    target = torch.randn_like(x)
147    t0 = time.perf_counter()
148    baseline_losses = [torch.mean((run_word(w, ops, x) - target) ** 2).item() for w in proposals]
149    baseline_time = time.perf_counter() - t0
150    t0 = time.perf_counter()
151    cache = {}
152    idea_losses = []
153    for w, key in zip(proposals, canonical_keys):
154        if key not in cache:
155            cache[key] = torch.mean((run_word(key, ops, x) - target) ** 2).item()
156        idea_losses.append(cache[key])
157    idea_time = time.perf_counter() - t0
158    max_loss_disagreement = max(abs(a - b) for a, b in zip(baseline_losses, idea_losses))
159
160    result = {
161        "seed": SEED,
162        "device": device,
163        "alphabet_size": len(names),
164        "word_length": 12,
165        "random_proposals": len(proposals),
166        "commuting_swap_trials": len(swap_errors),
167        "max_commuting_swap_error": max_commuting_error,
168        "median_noncommuting_swap_difference": median_noncommuting_difference,
169        "max_word_vs_canonical_error": max(canonical_errors),
170        "baseline_unique_strings": baseline_unique,
171        "canonical_unique_classes": idea_unique,
172        "deduplication_fraction": 1.0 - idea_unique / len(proposals),
173        "baseline_mean_loss": float(np.mean(baseline_losses)),
174        "idea_mean_loss": float(np.mean(idea_losses)),
175        "max_baseline_vs_cached_loss_disagreement": max_loss_disagreement,
176        "baseline_forward_seconds": baseline_time,
177        "idea_forward_seconds": idea_time,
178        "forward_reduction_fraction": 1.0 - idea_unique / len(proposals),
179        "cache_speedup": baseline_time / idea_time if idea_time else None,
180    }
181    print(json.dumps(result, indent=2))
182
183
184if __name__ == "__main__":
185    try:
186        main()
187    except Exception as exc:
188        # CUDA can fail in a shared environment; retrying explicitly on CPU
189        # keeps the experiment reproducible rather than silently fabricating it.
190        if torch.cuda.is_available():
191            print(json.dumps({"cuda_error": repr(exc), "retry": "cpu"}))
192            torch.cuda.is_available = lambda: False
193            main()
194        else:
195            raise