cleanup submission: model.py

model.py

@@ -66,6 +66,14 @@ class ModMulBP(ModularMultiplicationModel):
             self.regime = self.base ** self.W
         else:
             self.regime = 10 ** self.W
+        self.mulonly = bool(ckpt.get("mulonly", False))
+        if self.mulonly:
+            # Tier-0 member: pure multiplication. p never enters the trace --
+            # the operand width is the only constraint, so claim every p the
+            # specialists don't (router sorts by regime, this sorts last).
+            self.regime = 10 ** (2 * self.W)
+            self.p_lo, self.p_hi = 2, 2 ** 4096
+            return
         # Trained prime span (for router dispatch): derived from the ckpt's
         # tier list via the official tier geometry; fallback = full regime.
         self.p_lo, self.p_hi = 2, self.regime - 1
@@ -96,7 +104,13 @@ class ModMulBP(ModularMultiplicationModel):
     @torch.no_grad()
     def predict_digits_batch(self, inputs):
         out = [[0]] * len(inputs)
-        if self.composed:
+        if self.mulonly:
+            prompt_fn = lambda p, x, y, W: cenc.prompt_str(  # noqa: E731
+                x, y, p, W, 10, mulonly=True)
+            n_gen = cenc.gen_len(self.W, mulonly=True)
+            decode_fn = lambda g, W: cenc.decode_answer(     # noqa: E731
+                g, W, 10, mulonly=True)
+        elif self.composed:
             prompt_fn = lambda p, x, y, W: cenc.prompt_str(  # noqa: E731
                 x, y, p, W, self.base, self.subpad)
             n_gen = cenc.gen_len(self.W, self.scratch, self.cursor, self.subpad,
@@ -116,9 +130,12 @@ class ModMulBP(ModularMultiplicationModel):
                          else enc.decode_answer)
         prompts, idx = [], []
         for i, (a, b, p) in enumerate(inputs):
-            if p >= self.regime:          # outside trained width -> honest 0
-                continue
             x, y = a % p, b % p
+            if self.mulonly:
+                if max(x, y) >= 10 ** self.W:    # operands don't fit -> honest 0
+                    continue
+            elif p >= self.regime:        # outside trained width -> honest 0
+                continue
             prompts.append(prompt_fn(p, x, y, self.W))
             idx.append((i, x, y, p))
         if not prompts: