modmul router v0: t12 + composed-t3b members (pre-t3g)

14bef4a verified 5 months ago

3.89 kB

"""Shared sequence encoding for the modmul BP-install model.
The model learns the field-multiplication map  (p, x, y) -> (x*y) mod p  with
x, y in [0, p).  At inference the legal two-operand reduction a%p, b%p produces
x, y (the same step the baselines use); training samples x, y in [0, p) directly,
which matches the (a mod p) distribution for p << operand range.
Fixed-width, base-10, reverse-LSB answer (the proven arithmetic recipe). The
prime p is an explicit in-sequence conditioning field so one model generalises
across primes:
    p007*012,003=R120\n        # W=3 example: p=7, x=12%7=5? no -- x,y already <p
                               # e.g. p=7,x=5,y=4 -> 20%7=6 -> "p7*5,4=R6\n" (W=1)
Layout (width W = max decimal digits of any prime in scope):
    "p" P(W) "*" X(W) "," Y(W) "=R" REV_ANS(W) "\n"
zero-padded fields; REV_ANS is the W-digit answer reversed (least-significant
digit first), so digit k is emitted at a fixed position and the carry flows
left-to-right the way reverse-LSB addition does.
"""
from __future__ import annotations

MUL, SEP, EQ, REVMARK, NL = "*", ",", "=", "R", "\n"
NEWLINE_ID = ord("\n")


def width_for_primes(primes) -> int:
    """Decimal digits needed to hold the largest residue (p-1)."""
    return max(len(str(p - 1)) for p in primes)


def prompt_str(p: int, x: int, y: int, W: int) -> str:
    return f"p{p:0{W}d}*{x:0{W}d},{y:0{W}d}={REVMARK}"


def build_example(p: int, x: int, y: int, W: int):
    """Return (text, ann) where ann[t] holds the install label for the residual
    at byte position t (which predicts byte t+1), NTP-aligned."""
    ans = (x * y) % p
    rev_ans = f"{ans:0{W}d}"[::-1]            # LSB-first, width W
    prompt = prompt_str(p, x, y, W)
    text = prompt + rev_ans + NL
    ann = [dict() for _ in range(len(text))]
    base = len(prompt)                        # first answer char index
    for k in range(W):                        # answer digit k (LSB-first)
        pos = base + k
        ann[pos - 1]["ans"] = int(rev_ans[k])  # residual at pos-1 predicts it
    return text, ann


def answer_len(W: int) -> int:
    return W


# --- Scratchpad ("chain-of-thought") inference variant -------------------
# The model GENERATES the product as reverse-LSB digits before the answer:
#     "p" P(W) "*" X(W) "," Y(W) "=" REV_PROD(2W) "R" REV_ANS(W) "\n"
# At inference we feed up to '=' and greedily decode the full scratchpad; the
# answer is the last W generated digits. Python never computes x*y -- the model
# emits the product digits (a learned circuit). Only a%p, b%p is done in code.

def prompt_str_sp(p: int, x: int, y: int, W: int) -> str:
    return f"p{p:0{W}d}*{x:0{W}d},{y:0{W}d}{EQ}"


def scratchpad_len(W: int) -> int:
    return 2 * W + 1 + W


def decode_answer_sp(gen_chars: str, W: int) -> int:
    """Generated suffix is 2W product digits, 'R', then W answer digits."""
    return decode_answer(gen_chars[2 * W + 1 : 2 * W + 1 + W], W)


# --- Schoolbook ("partial products") inference variant -------------------
# Generated suffix: W partial products (W+1 wide), 'P', 2W product digits,
# 'R', W answer digits. The model emits the partial products and sums them
# (learned circuit); the answer is the final W digits.

def school_gen_len(W: int) -> int:
    return W * (W + 1) + 1 + 2 * W + 1 + W


def decode_answer_school(gen_chars: str, W: int) -> int:
    start = W * (W + 1) + 1 + 2 * W + 1
    return decode_answer(gen_chars[start:start + W], W)


def decode_answer(gen_chars: str, W: int) -> int:
    """Greedy-generated W chars are the reverse-LSB answer digits. Non-digits
    (untrained junk) count as 0. Returns the integer answer."""
    val = 0
    for k in range(min(W, len(gen_chars))):
        c = gen_chars[k]
        d = int(c) if "0" <= c <= "9" else 0  # ASCII only ('³'.isdigit() is True but int() fails)
        val += d * (10 ** k)
    return val

	"""Shared sequence encoding for the modmul BP-install model.

	The model learns the field-multiplication map (p, x, y) -> (x*y) mod p with
	x, y in [0, p). At inference the legal two-operand reduction a%p, b%p produces
	x, y (the same step the baselines use); training samples x, y in [0, p) directly,
	which matches the (a mod p) distribution for p << operand range.

	Fixed-width, base-10, reverse-LSB answer (the proven arithmetic recipe). The
	prime p is an explicit in-sequence conditioning field so one model generalises
	across primes:

	p007*012,003=R120\n # W=3 example: p=7, x=12%7=5? no -- x,y already <p
	# e.g. p=7,x=5,y=4 -> 20%7=6 -> "p7*5,4=R6\n" (W=1)

	Layout (width W = max decimal digits of any prime in scope):
	"p" P(W) "*" X(W) "," Y(W) "=R" REV_ANS(W) "\n"
	zero-padded fields; REV_ANS is the W-digit answer reversed (least-significant
	digit first), so digit k is emitted at a fixed position and the carry flows
	left-to-right the way reverse-LSB addition does.
	"""
	from __future__ import annotations

	MUL, SEP, EQ, REVMARK, NL = "*", ",", "=", "R", "\n"
	NEWLINE_ID = ord("\n")


	def width_for_primes(primes) -> int:
	"""Decimal digits needed to hold the largest residue (p-1)."""
	return max(len(str(p - 1)) for p in primes)


	def prompt_str(p: int, x: int, y: int, W: int) -> str:
	return f"p{p:0{W}d}*{x:0{W}d},{y:0{W}d}={REVMARK}"


	def build_example(p: int, x: int, y: int, W: int):
	"""Return (text, ann) where ann[t] holds the install label for the residual
	at byte position t (which predicts byte t+1), NTP-aligned."""
	ans = (x * y) % p
	rev_ans = f"{ans:0{W}d}"[::-1] # LSB-first, width W
	prompt = prompt_str(p, x, y, W)
	text = prompt + rev_ans + NL
	ann = [dict() for _ in range(len(text))]
	base = len(prompt) # first answer char index
	for k in range(W): # answer digit k (LSB-first)
	pos = base + k
	ann[pos - 1]["ans"] = int(rev_ans[k]) # residual at pos-1 predicts it
	return text, ann


	def answer_len(W: int) -> int:
	return W


	# --- Scratchpad ("chain-of-thought") inference variant -------------------
	# The model GENERATES the product as reverse-LSB digits before the answer:
	# "p" P(W) "*" X(W) "," Y(W) "=" REV_PROD(2W) "R" REV_ANS(W) "\n"
	# At inference we feed up to '=' and greedily decode the full scratchpad; the
	# answer is the last W generated digits. Python never computes x*y -- the model
	# emits the product digits (a learned circuit). Only a%p, b%p is done in code.

	def prompt_str_sp(p: int, x: int, y: int, W: int) -> str:
	return f"p{p:0{W}d}*{x:0{W}d},{y:0{W}d}{EQ}"


	def scratchpad_len(W: int) -> int:
	return 2 * W + 1 + W


	def decode_answer_sp(gen_chars: str, W: int) -> int:
	"""Generated suffix is 2W product digits, 'R', then W answer digits."""
	return decode_answer(gen_chars[2 * W + 1 : 2 * W + 1 + W], W)


	# --- Schoolbook ("partial products") inference variant -------------------
	# Generated suffix: W partial products (W+1 wide), 'P', 2W product digits,
	# 'R', W answer digits. The model emits the partial products and sums them
	# (learned circuit); the answer is the final W digits.

	def school_gen_len(W: int) -> int:
	return W * (W + 1) + 1 + 2 * W + 1 + W


	def decode_answer_school(gen_chars: str, W: int) -> int:
	start = W * (W + 1) + 1 + 2 * W + 1
	return decode_answer(gen_chars[start:start + W], W)


	def decode_answer(gen_chars: str, W: int) -> int:
	"""Greedy-generated W chars are the reverse-LSB answer digits. Non-digits
	(untrained junk) count as 0. Returns the integer answer."""
	val = 0
	for k in range(min(W, len(gen_chars))):
	c = gen_chars[k]
	d = int(c) if "0" <= c <= "9" else 0 # ASCII only ('³'.isdigit() is True but int() fails)
	val += d * (10 ** k)
	return val