cleanup submission: composed_encoding.py

composed_encoding.py

@@ -37,11 +37,14 @@ AMARK, BMARK, DIVMARK, EQ, ACC, NMARK, REVMARK, NL = "A", "B", "M", "=", "a", "N
 PLSB = "m"
 
 
-def prompt_str(x: int, y: int, p: int, W: int, base: int, subpad: bool = False) -> str:
+def prompt_str(x: int, y: int, p: int, W: int, base: int, subpad: bool = False,
+               mulonly: bool = False) -> str:
     """subpad adds p LSB-first to the prompt: the MSB field serves quotient
     estimation, the LSB field aligns with the LSB-emitted qd*p / remainder
     chains (per-step diagnostic: subtraction digits need REVERSED access into
     the MSB-only field — the aligned acc stage learned, the division didn't)."""
+    if mulonly:                                  # tier 0: pure multiplication
+        return AMARK + limb_str(x, W, base) + BMARK + limb_str(y, W, base) + EQ
     s = (AMARK + limb_str(x, W, base) + BMARK + limb_str(y, W, base)
          + DIVMARK + limb_str(p, W, base, msb_first=True))
     if subpad:
@@ -52,7 +55,8 @@ def prompt_str(x: int, y: int, p: int, W: int, base: int, subpad: bool = False)
 def gen_len(W: int, scratch: bool = False, cursor: bool = False,
             subpad: bool = False, stepidx: bool = False,
             skiptriv: bool = False, subnum: bool = False,
-            bemit: bool = False, srt: bool = False) -> int:
+            bemit: bool = False, srt: bool = False,
+            mulonly: bool = False) -> int:
     """Tokens generated after '=': W acc blocks, N restage, quotient, answer.
     scratch=True: each quotient limb is followed by the running remainder
     (W limbs, LSB-first) — externalizes the division state, which the model
@@ -100,6 +104,8 @@ def gen_len(W: int, scratch: bool = False, cursor: bool = False,
     rem errors were 58% unstructured noise, 0.54 even at b=0 -- the
     install-only borrow supervision never built a circuit; emit the chain
     (the scratch lesson, applied to the borrow state)."""
+    if mulonly:                                  # acc rows + answer copy
+        return W * (1 + 2 * W) + 1 + 2 * W
     restage = 1 + (3 if stepidx else 1) * 2 * W
     pcopy = (1 + W) if skiptriv else 0
     if srt:
@@ -144,8 +150,9 @@ def _srt_qd(num: int, p: int, base: int) -> int:
 def build_example(x: int, y: int, p: int, W: int, base: int, scratch: bool = False,
                   cursor: bool = False, subpad: bool = False, stepidx: bool = False,
                   skiptriv: bool = False, subnum: bool = False,
-                  bemit: bool = False, srt: bool = False):
-    """Return (text, ann). x,y < p < base**W."""
+                  bemit: bool = False, srt: bool = False, mulonly: bool = False,
+                  srt_dither: float = 0.0):
+    """Return (text, ann). x,y < p < base**W (mulonly: x,y < base**W)."""
     assert not subnum or subpad, "subnum requires subpad"
     assert not bemit or subnum, "bemit requires subnum"
     assert not srt or (skiptriv and bemit), "srt requires skiptriv+bemit"
@@ -178,6 +185,22 @@ def build_example(x: int, y: int, p: int, W: int, base: int, scratch: bool = Fal
         pos += 2 * W
     assert acc == N
 
+    if mulonly:                                  # tier 0: answer = N itself
+        parts.append(NMARK)
+        pos += 1
+        n_l = to_limbs(N, 2 * W, base)
+        for k in range(2 * W):
+            ann_g.append((pos + k, "ans", n_l[k]))
+        parts.append(limb_str(N, 2 * W, base) + NL)
+        pos += 2 * W + 1
+        prompt = prompt_str(x, y, p, W, base, mulonly=True)
+        text = prompt + "".join(parts)
+        assert len(text) == len(prompt) + gen_len(W, mulonly=True) + 1
+        ann = [dict() for _ in range(len(text))]
+        for gi, var, val in ann_g:
+            ann[len(prompt) + gi - 1][var] = int(val)
+        return text, ann
+
     n_msb = to_limbs(N, 2 * W, base, msb_first=True)
     if stepidx:                                  # restage as (index, limb) pairs
         parts.append(NMARK + "".join(f"{i:02d}" + limb_char(n_msb[i])
@@ -212,6 +235,18 @@ def build_example(x: int, y: int, p: int, W: int, base: int, scratch: bool = Fal
             parts.append("".join(limb_char(v) for v in num_c))
             pos += K
             qd = _srt_qd(num, p, base)
+            if srt_dither:
+                # Deterministic pseudo-random alternate-digit training: SRT's
+                # redundancy makes qd+-1 often equally valid, but the model
+                # only ever saw label-digit continuations -- free-running it
+                # sometimes picks the OTHER valid digit and must continue
+                # consistently from its own choice. Dithered traces teach
+                # exactly those continuations. (hash-derived: reproducible)
+                h = hash((x, y, p, i)) & 0xFFFF
+                if h < int(srt_dither * 0x10000):
+                    alt = qd + (1 if h & 1 else -1)
+                    if abs(alt) <= base - 1 and abs(num - alt * p) < p:
+                        qd = alt
             assert abs(num - qd * p) < p, (x, y, p, i, num, qd)
             ann_g.append((pos, "qhat", qd + base - 1))   # 19-class signed label
             parts.append(chr(0x90 + base - 1 + qd))
@@ -331,7 +366,9 @@ def build_example(x: int, y: int, p: int, W: int, base: int, scratch: bool = Fal
 
 
 def var_specs(W: int, base: int, scratch: bool = False, subpad: bool = False,
-              srt: bool = False):
+              srt: bool = False, mulonly: bool = False):
+    if mulonly:                                  # tier 0: multiply only
+        return [("ans", base), ("acarry", base)]
     specs = [("ans", base), ("qhat", 2 * base - 1 if srt else base),
              ("acarry", base)]
     if subpad:                                  # multiply-carry + subtract-borrow chains
@@ -345,7 +382,10 @@ def decode_answer(gen_chars: str, W: int, base: int, scratch: bool = False,
                   cursor: bool = False, subpad: bool = False,
                   stepidx: bool = False, skiptriv: bool = False,
                   subnum: bool = False, bemit: bool = False,
-                  srt: bool = False) -> int:
+                  srt: bool = False, mulonly: bool = False) -> int:
+    if mulonly:                                  # answer = 2W limbs after 'N'
+        off = gen_len(W, mulonly=True) - 2 * W
+        return parse_limbs(gen_chars[off:off + 2 * W], base)
     # the answer block after the R-mark is W plain limbs in every mode
     off = gen_len(W, scratch, cursor, subpad, stepidx, skiptriv, subnum, bemit,
                   srt) - W