Single file Space, app.py self installs deps, CNN and 1 head attention proofs

app.py

@@ -1,298 +1,275 @@
+# app.py, single file Space
+# FieldSpace, Prime-Only Machine, mod 2^K exactness proofs
+import sys, subprocess, time, math, json
+
+def _ensure(package_spec: str):
+    try:
+        __import__(package_spec.split("==")[0].split(">=")[0].split("<")[0])
+        return
+    except Exception:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", package_spec])
+
+# Self install, single file repo, no external requirements
+for spec in ["gradio==4.44.1", "numpy>=1.26", "torch>=2.2,<2.4"]:
+    _ensure(spec)
 
-import gradio as gr
 import numpy as np
+import torch
+import gradio as gr
+
+torch.set_grad_enabled(False)
+torch.set_num_threads(1)
+DEVICE = "cpu"
 
 # -----------------------------
-# Utilities: RNS + CRT
+# Utilities (mod 2^K)
 # -----------------------------
-def pick_moduli(max_abs, primes=(257,263,269,271,277,281)):
-    # choose K so 2^(K-1) > max_abs
-    K = max(20, int(max_abs).bit_length()+1)
-    MOD2 = 1 << K
-    # ensure product > 2*max_abs (with plenty margin for toy sizes this is fine)
-    M = MOD2
-    ps = []
-    for p in primes:
-        ps.append(p); M *= p
-        if M > 2*max_abs: break
-    return K, MOD2, tuple(ps)
-
-def encode_rns(x, MOD2, primes):
-    outs = [np.asarray(x, dtype=object) & (MOD2-1)]
-    for p in primes:
-        outs.append(np.asarray(x, dtype=object) % p)
-    return np.stack(outs, axis=0)  # [L, *shape]
-
-def batched_crt(res, MOD2, primes):
-    # res: [L, *shape], L=1+len(primes), dtype=object (python ints)
-    moduli = (MOD2,)+tuple(primes)
-    L = res.shape[0]
-    M = 1
-    for m in moduli: M *= m
-    Mi = [M//m for m in moduli]
-    inv = [pow(int(Mi[i] % moduli[i]), -1, moduli[i]) for i in range(L)]
-    flat = res.reshape(L, -1)
-    out = []
-    half = M//2
-    for k in range(flat.shape[1]):
-        acc = 0
-        for i in range(L):
-            acc = (acc + (flat[i,k] * inv[i] * Mi[i])) % M
-        if acc > half: acc -= M
-        out.append(int(acc))
-    return np.array(out, dtype=np.int64).reshape(res.shape[1:])
+def seed_all(s: int):
+    return torch.Generator(device=DEVICE).manual_seed(int(s))
+
+def choose_K(max_abs: int, min_bits: int = 24) -> int:
+    need = max(min_bits, int(max_abs).bit_length() + 1)
+    return min(62, max(need, min_bits))
+
+def mod2_encode(x: torch.Tensor, MOD2: int) -> torch.Tensor:
+    return (x & (MOD2 - 1)).to(torch.int64)
+
+def mod2_center(v: torch.Tensor, MOD2: int) -> torch.Tensor:
+    K = MOD2.bit_length() - 1
+    half = 1 << (K - 1)
+    out = v.clone()
+    out[out > half] -= MOD2
+    return out
 
 # -----------------------------
-# CNN proof (tiny shapes)
+# CNN (Conv -> Act -> Conv -> Act -> FC)
 # -----------------------------
 def im2col_nchw_int(x, kH, kW, stride=1, pad=0):
-    N,C,H,W = x.shape
-    Hp = H+2*pad; Wp = W+2*pad
-    xp = np.zeros((N,C,Hp,Wp), dtype=np.int64)
-    xp[:,:,pad:pad+H, pad:pad+W] = x
-    Hout = (Hp - kH)//stride + 1
-    Wout = (Wp - kW)//stride + 1
+    N, C, H, W = x.shape
+    Hout = (H + 2*pad - kH)//stride + 1
+    Wout = (W + 2*pad - kW)//stride + 1
+    if pad:
+        x_pad = torch.zeros((N, C, H + 2*pad, W + 2*pad), dtype=x.dtype, device=x.device)
+        x_pad[:, :, pad:pad+H, pad:pad+W] = x
+    else:
+        x_pad = x
     cols = []
     for i in range(Hout):
         for j in range(Wout):
-            patch = xp[:,:,i*stride:i*stride+kH, j*stride:j*stride+kW]
-            cols.append(patch.reshape(N,-1))
-    out = np.stack(cols, axis=1)   # [N, Hout*Wout, C*kH*kW]
-    out = out.reshape(N*Hout*Wout, -1)
+            patch = x_pad[:, :, i*stride:i*stride+kH, j*stride:j*stride+kW]
+            cols.append(patch.reshape(N, -1))
+    out = torch.stack(cols, dim=1)           # [N, Hout*Wout, C*kH*kW]
+    out = out.reshape(N*Hout*Wout, -1)       # [N*Hout*Wout, C*kH*kW]
     return out, Hout, Wout
 
-def conv_ref(X, W, kH, kW, stride=1, pad=0):
-    N,C,H,W = X.shape
-    Cout,Cin,KH,KW = W.shape
-    Xcol,Hout,Wout = im2col_nchw_int(X,kH,kW,stride,pad)
-    Ycol = Xcol @ W.reshape(Cout,-1).T
-    Y = Ycol.reshape(N,Hout,Wout,Cout).transpose(0,3,1,2).copy()
-    return Y
-
-def relu_ref(x): return np.where(x<0, 0, x)
-def poly_ref(x, s): return (x*x) >> s
-def linear_ref(X, W): return X @ W
-
-def relu_rns(Xr, MOD2, primes):
-    # Xr: [L, *shape]
-    L = Xr.shape[0]
-    K = MOD2.bit_length()-1
-    x2 = Xr[0]
-    mask = (x2 < (1<<(K-1))).astype(np.int64)  # 1 if non-negative, else 0
-    out = []
-    out.append((x2 * mask) & (MOD2-1))
-    for i,p in enumerate(primes, start=1):
-        out.append((Xr[i] * mask) % p)
-    return np.stack(out, axis=0)
-
-def poly_rns(Xr, s, MOD2, primes):
-    # q = floor(x^2 / 2^s)
-    mask_s = (1<<s)-1
-    out = []
-    # 2^K channel
-    y2 = (Xr[0]*Xr[0]) & (MOD2-1)
-    q2 = (y2 >> s) & (MOD2-1)
-    out.append(q2)
-    # r from low s bits of x only
-    low = Xr[0] & mask_s
-    r = (low*low) & mask_s
-    for i,p in enumerate(primes, start=1):
-        y2p = (Xr[i]*Xr[i]) % p
-        inv = pow(2, -s, p)
-        qi = ((y2p - (r % p)) * inv) % p
-        out.append(qi)
-    return np.stack(out, axis=0)
-
-def matmul_rns(Ar, Br, MOD2, primes):
-    # Ar: [L, M, K], Br: [L, K, N] -> [L, M, N]
-    outs=[]
-    outs.append((Ar[0] @ Br[0]) & (MOD2-1))
-    for i,p in enumerate(primes, start=1):
-        outs.append((Ar[i] @ Br[i]) % p)
-    return np.stack(outs, axis=0)
-
-def conv_rns(Xr, Wr, kH, kW, stride=1, pad=0, MOD2=0, primes=()):
-    # reshape per-modulus and reuse im2col on ints
-    L = Xr.shape[0]
-    N,C,H,W = Xr[0].shape
-    Cout = Wr[0].shape[0]
-    Ymods=[]
-    for i in range(L):
-        Xi = Xr[i].astype(object)
-        Wi = Wr[i].astype(object)
-        Xcol,Hout,Wout = im2col_nchw_int(np.asarray(Xi, dtype=np.int64), kH,kW,stride,pad)
-        Wcol = np.asarray(Wi.reshape(Cout,-1), dtype=np.int64)
-        if i==0:
-            MOD = MOD2
-            Ycol = (Xcol @ Wcol.T) & (MOD-1)
-        else:
-            p = primes[i-1]
-            Ycol = (Xcol @ Wcol.T) % p
-        Y = Ycol.reshape(N,Hout,Wout,Cout).transpose(0,3,1,2).copy()
-        Ymods.append(Y.astype(object))
-    return np.stack(Ymods, axis=0)
-
-def cnn_proof(seed, xmax, wmax, activation, shift):
-    # Tiny shapes to keep it snappy
-    B,Cin,H,W = 2,1,8,8
-    Cout = 2
-    KH=KW=3
-    CLS=8
-    rng = np.random.default_rng(seed)
-    X   = rng.integers(-xmax, xmax+1, size=(B,Cin,H,W), dtype=np.int64)
-    Wc  = rng.integers(-wmax, wmax+1, size=(Cout,Cin,KH,KW), dtype=np.int64)
-    Wfc = rng.integers(-wmax, wmax+1, size=(Cout*H*W, CLS), dtype=np.int64)
-
-    # Reference (int64)
-    Y = conv_ref(X,Wc,KH,KW,1,1)
-    A = relu_ref(Y) if activation=="relu" else poly_ref(Y, shift)
-    Z = linear_ref(A.reshape(B,-1), Wfc)
-
-    # Bounds (very conservative)
-    b1 = int(Cin*KH*KW * xmax * wmax)
-    a1 = b1 if activation=="relu" else (b1*b1)>>shift
-    b2 = int(Cout*KH*KW * a1 * wmax)
-    a2 = b2 if activation=="relu" else (b2*b2)>>shift
-    fc = int((Cout*H*W) * a2 * wmax)
-    MAX_ABS = max(abs(b1),abs(a1),abs(b2),abs(a2),abs(fc), 1)
-
-    K, MOD2, primes = pick_moduli(max_abs=MAX_ABS)
-    # Encode
-    Xr  = encode_rns(X, MOD2, primes)
-    Wcr = encode_rns(Wc, MOD2, primes)
-    Wfr = encode_rns(Wfc, MOD2, primes)
-
-    # RNS path
-    Yr = conv_rns(Xr,Wcr,KH,KW,1,1,MOD2,primes)
-    Ar = relu_rns(Yr, MOD2, primes) if activation=="relu" else poly_rns(Yr, shift, MOD2, primes)
-    Arf = Ar.reshape(Ar.shape[0], B, -1)
-    Zr  = matmul_rns(Arf, Wfr, MOD2, primes)
-    Zrec= batched_crt(Zr, MOD2, primes)
-
-    ok = np.array_equal(Z, Zrec)
-    arg = np.array_equal(Z.argmax(1), Zrec.argmax(1))
-
-    log = []
-    log.append("Stage-by-stage equality (exact):")
-    log.append(f"conv : {'✅' if np.array_equal(batched_crt(Yr,MOD2,primes), Y) else '❌'}")
-    log.append(f"act  : {'✅' if np.array_equal(batched_crt(Ar,MOD2,primes), A) else '❌'}")
-    log.append(f"fc   : {'✅' if ok else '❌'}")
-    log.append(f"Argmax match: {arg} (ref={Z.argmax(1).tolist()}, prime={Zrec.argmax(1).tolist()})")
+def conv_int64_im2col(X, W, kH, kW, stride=1, pad=0):
+    Xcol, Hout, Wout = im2col_nchw_int(X, kH, kW, stride=stride, pad=pad)
+    Ycol = Xcol @ W.reshape(W.shape[0], -1).t()
+    return Ycol.reshape(X.shape[0], Hout, Wout, W.shape[0]).permute(0,3,1,2).contiguous()
+
+def run_cnn_proof(seed:int=0, Xmax:int=31, Wmax:int=31, act_kind:str="relu", poly_shift:int=7):
+    g = seed_all(seed)
+    B,Cin,H,W  = 4,1,16,16
+    C1,C2      = 8,8
+    kH,kW      = 3,3
+    STRIDE,PAD = 1,1
+    CLS        = 10
+
+    X  = torch.randint(-Xmax, Xmax+1, (B,Cin,H,W), dtype=torch.int64, generator=g, device=DEVICE)
+    W1 = torch.randint(-Wmax, Wmax+1, (C1,Cin,kH,kW), dtype=torch.int64, generator=g, device=DEVICE)
+    W2 = torch.randint(-Wmax, Wmax+1, (C2,C1,kH,kW), dtype=torch.int64, generator=g, device=DEVICE)
+    Wf = torch.randint(-Wmax, Wmax+1, (C2*H*W, CLS),   dtype=torch.int64, generator=g, device=DEVICE)
+
+    t0 = time.time()
+    Z1 = conv_int64_im2col(X, W1, kH,kW, STRIDE, PAD)
+    if act_kind=="relu":
+        A1 = Z1.clamp_min_(0)
+    else:
+        A1 = ((Z1*Z1) >> int(poly_shift))
+    Z2 = conv_int64_im2col(A1, W2, kH,kW, STRIDE, PAD)
+    if act_kind=="relu":
+        A2 = Z2.clamp_min_(0)
+    else:
+        A2 = ((Z2*Z2) >> int(poly_shift))
+    Y  = (A2.reshape(B, -1) @ Wf)
+    t1 = time.time()
+
+    max_abs = int(max(Z1.abs().max().item(), A1.abs().max().item(),
+                      Z2.abs().max().item(), A2.abs().max().item(), Y.abs().max().item()))
+    K    = choose_K(max_abs, min_bits=24)
+    MOD2 = 1 << K
 
-    report = {
-        "activation": activation,
-        "shift": int(shift),
-        "K": int(K),
-        "primes": list(map(int,primes)),
-        "exact": bool(ok),
-        "argmax_exact": bool(arg),
-        "ref_logits_sample": Z[0,:min(8,CLS)].tolist(),
-        "rec_logits_sample": Zrec[0,:min(8,CLS)].tolist(),
+    Z1m = mod2_encode(conv_int64_im2col(mod2_encode(X,MOD2), mod2_encode(W1,MOD2), kH,kW, STRIDE, PAD), MOD2)
+    Z1i = mod2_center(Z1m, MOD2)
+    A1i = Z1i.clamp_min_(0) if act_kind=="relu" else ((Z1i*Z1i) >> int(poly_shift))
+
+    Z2m = mod2_encode(conv_int64_im2col(mod2_encode(A1i,MOD2), mod2_encode(W2,MOD2), kH,kW, STRIDE, PAD), MOD2)
+    Z2i = mod2_center(Z2m, MOD2)
+    A2i = Z2i.clamp_min_(0) if act_kind=="relu" else ((Z2i*Z2i) >> int(poly_shift))
+
+    Ym  = mod2_encode((mod2_encode(A2i.reshape(B,-1),MOD2) @ mod2_encode(Wf,MOD2)), MOD2)
+    Yi  = mod2_center(Ym, MOD2)
+
+    ok_all = bool(torch.equal(Y, Yi))
+    arg_ok = bool(torch.equal(Y.argmax(1), Yi.argmax(1)))
+
+    txt = []
+    txt.append(f"K={K} (2^(K-1)={1<<(K-1):,} > max_abs={max_abs:,}), act={act_kind} ({'ReLU' if act_kind=='relu' else f'poly SHIFT={poly_shift}'})")
+    txt.append("Stage by stage equality, exact:")
+    txt.append(f"Z1: {'OK' if torch.equal(Z1,Z1i) else 'NO'} | A1: {'OK' if torch.equal(A1,A1i) else 'NO'} | "
+               f"Z2: {'OK' if torch.equal(Z2,Z2i) else 'NO'} | Y: {'OK' if torch.equal(Y,Yi) else 'NO'}")
+    txt.append(f"Argmax match: {arg_ok} (ref={Y.argmax(1).tolist()}, mod2={Yi.argmax(1).tolist()})")
+    txt.append(f"Timing ms: ref≈{(t1-t0)*1000:.2f}")
+    return "\n".join(txt), {
+        "ok_all": ok_all, "argmax_ok": arg_ok, "K": K,
+        "act": act_kind, "poly_shift": int(poly_shift),
+        "ref_top1": Y.argmax(1).tolist(), "mod2_top1": Yi.argmax(1).tolist()
     }
-    return "\n".join(log), report
 
 # -----------------------------
-# Attention (exact numerators)
+# 1 head Attention
 # -----------------------------
-def attn_numerators_proof(seed):
-    rng = np.random.default_rng(seed)
-    T,d = 2,4
-    xmax = 15
-    t = 7  # scale shift
-
-    Q = rng.integers(-xmax,xmax+1,size=(T,d),dtype=np.int64)
-    K_ = rng.integers(-xmax,xmax+1,size=(T,d),dtype=np.int64)
-    V = rng.integers(-xmax,xmax+1,size=(T,d),dtype=np.int64)
-
-    Sraw = Q @ K_.T
-    S = Sraw >> t   # exact floor
-    N = S @ V       # numerators
-
-    # conservative bound
-    s_max = int((d * xmax * xmax) >> t)
-    n_max = int(T * s_max * xmax)
-    Kpow, MOD2, primes = pick_moduli(max_abs=max(abs(s_max),abs(n_max),1))
-
-    Qr  = encode_rns(Q,MOD2,primes)
-    Kr  = encode_rns(K_,MOD2,primes)
-    Vr  = encode_rns(V,MOD2,primes)
-
-    # Sraw_r and S_r
-    Sraw_r = matmul_rns(Qr, Kr.transpose(0,2,1), MOD2, primes)
-    # floor by right shift in 2^K channel + multiplicative inverse in odd primes
-    # (here scaling is power-of-two so we can just shift in 2^K, and multiply by inv(2^t) after peeling low bits = 0)
-    # since we already floored by pure shift in reference, it's safe to do:
-    Sr = []
-    Sr.append( (Sraw_r[0] >> t) & (MOD2-1) )
-    invs = [pow(2,-t,p) for p in primes]
-    for i,p in enumerate(primes, start=1):
-        Sr.append( (Sraw_r[i] * invs[i-1]) % p )
-    Sr = np.stack(Sr, axis=0)
-
-    Nr = matmul_rns(Sr, Vr, MOD2, primes)
-    Srec = batched_crt(Sr, MOD2, primes)
-    Nrec = batched_crt(Nr, MOD2, primes)
-
-    okS = np.array_equal(S, Srec)
-    okN = np.array_equal(N, Nrec)
-
-    log = []
-    log.append("Attention (exact numerators):")
-    log.append(f"S  : {'✅' if okS else '❌'}")
-    log.append(f"SV : {'✅' if okN else '❌'}")
+def split_heads(Z, B,T,C,H):
+    Dh = C//H
+    return Z.reshape(B,T,H,Dh).permute(0,2,1,3).contiguous()
+
+def merge_heads(Zh, B,T,C,H):
+    return Zh.permute(0,2,1,3).reshape(B,T,C).contiguous()
+
+def pick_shift_for_cap(smax:int, cap:int=8) -> int:
+    if smax <= cap: return 0
+    s = 0
+    while (smax >> s) > cap: s += 1
+    return s
+
+def run_attn_proof(seed:int=0, Xmax:int=7, Wmax:int=7, B:int=1, T:int=8, C:int=16, H:int=1):
+    assert C % H == 0
+    Dh = C//H
+    g  = seed_all(seed)
+
+    X  = torch.randint(0, Xmax+1, (B,T,C), dtype=torch.int64, generator=g, device=DEVICE)
+    Wq = torch.randint(0, Wmax+1, (C,C), dtype=torch.int64, generator=g, device=DEVICE)
+    Wk = torch.randint(0, Wmax+1, (C,C), dtype=torch.int64, generator=g, device=DEVICE)
+    Wv = torch.randint(0, Wmax+1, (C,C), dtype=torch.int64, generator=g, device=DEVICE)
+    Wo = torch.randint(0, Wmax+1, (C,C), dtype=torch.int64, generator=g, device=DEVICE)
+
+    t0 = time.time()
+    Q = X @ Wq; K = X @ Wk; V = X @ Wv
+    Qh, Kh, Vh = split_heads(Q,B,T,C,H), split_heads(K,B,T,C,H), split_heads(V,B,T,C,H)
+    Sraw = torch.einsum("bhtd,bhTd->bhtT", Qh, Kh)
+    SHIFT = pick_shift_for_cap(int(Sraw.max().item()), cap=8)
+    S = (Sraw >> SHIFT).clamp_min_(0)
+    E = (torch.ones_like(S) << S)
+    Den = E.sum(-1, keepdim=True)
+    Num = torch.einsum("bhtT,bhTd->bhtd", E, Vh)
+    Oh  = Num // Den
+    O   = merge_heads(Oh,B,T,C,H)
+    Y   = O @ Wo
+    t1 = time.time()
+
+    max_abs = int(max(Q.abs().max().item(), K.abs().max().item(), V.abs().max().item(),
+                      Sraw.abs().max().item(), S.abs().max().item(),
+                      Num.abs().max().item(), O.abs().max().item(), Y.abs().max().item()))
+    Kbits = choose_K(max_abs, min_bits=24)
+    MOD2  = 1 << Kbits
+
+    Qm = mod2_encode(X,MOD2) @ mod2_encode(Wq,MOD2); Qi = mod2_center(Qm,MOD2)
+    Km = mod2_encode(X,MOD2) @ mod2_encode(Wk,MOD2); Ki = mod2_center(Km,MOD2)
+    Vm = mod2_encode(X,MOD2) @ mod2_encode(Wv,MOD2); Vi = mod2_center(Vm,MOD2)
+    assert torch.equal(Qi,Q) and torch.equal(Ki,K) and torch.equal(Vi,V)
+
+    Qh, Kh, Vh = split_heads(Qi,B,T,C,H), split_heads(Ki,B,T,C,H), split_heads(Vi,B,T,C,H)
+
+    BH = B*H
+    Sraw_mod = torch.empty((BH, T, T), dtype=torch.int64)
+    Q2, K2 = Qh.reshape(BH,T,C//H), Kh.reshape(BH,T,C//H)
+    for b in range(BH):
+        Sraw_mod[b] = (mod2_encode(Q2[b],MOD2) @ mod2_encode(K2[b].t(),MOD2)) & (MOD2-1)
+    Sraw_i = mod2_center(Sraw_mod.reshape(B,H,T,T), MOD2)
+
+    if SHIFT == 0:
+        S_i = Sraw_i
+    else:
+        S2   = mod2_encode(Sraw_i, MOD2)
+        q2   = (S2 >> SHIFT) & (MOD2 - 1)
+        S_i  = mod2_center(q2, MOD2)
+    S_i = S_i.clamp_min_(0)
+
+    E_mod = (torch.ones_like(S_i) << S_i) & (MOD2 - 1)
+    Den_m = E_mod.sum(-1)
+    Den_i = mod2_center(Den_m, MOD2)
+
+    Vh_m = mod2_encode(Vh, MOD2)
+    Num_m = torch.empty((B,H,T, C//H), dtype=torch.int64)
+    for b in range(B):
+        for h in range(H):
+            Num_m[b,h] = (E_mod[b,h] @ Vh_m[b,h]) & (MOD2 - 1)
+    Num_i = mod2_center(Num_m, MOD2)
+
+    Oh = Num_i // Den_i.unsqueeze(-1)
+    Oi = merge_heads(Oh,B,T,C,H)
+
+    Ym = (mod2_encode(Oi,MOD2) @ mod2_encode(Wo,MOD2)) & (MOD2 - 1)
+    Yi = mod2_center(Ym, MOD2)
+
+    ok_all = bool(torch.equal(Y, Yi))
+    arg_ok = bool(torch.equal(Y.argmax(-1), Yi.argmax(-1)))
+
+    lines = []
+    lines.append(f"Auto SHIFT={SHIFT} cap=8")
+    lines.append(f"K={Kbits} 2^(K-1)={1<<(Kbits-1):,} > max_abs={max_abs:,}")
+    lines.append("Stage by stage equality, exact:")
+    lines.append("Q:{0} K:{1} V:{2} S:{3} Y:{4}".format(
+        'OK' if torch.equal(Q,Qi) else 'NO',
+        'OK' if torch.equal(K,Ki) else 'NO',
+        'OK' if torch.equal(V,Vi) else 'NO',
+        'OK' if torch.equal(S,S_i) else 'NO',
+        'OK' if torch.equal(Y,Yi) else 'NO'
+    ))
+    lines.append(f"Argmax match: {arg_ok} ref={Y.argmax(-1).tolist()} mod2={Yi.argmax(-1).tolist()}")
+    lines.append(f"Timing ms: ref≈{(t1-t0)*1000:.2f}")
 
     report = {
-        "K": int(Kpow),
-        "primes": list(map(int,primes)),
-        "S_exact": bool(okS),
-        "N_exact": bool(okN),
-        "S_sample": S.tolist(),
-        "Srec_sample": Srec.tolist(),
-        "N_sample": N.tolist(),
-        "Nrec_sample": Nrec.tolist(),
+        "ok_all": ok_all, "argmax_ok": arg_ok, "K": Kbits,
+        "ref_top1": Y.argmax(-1).tolist(), "mod2_top1": Yi.argmax(-1).tolist()
     }
-    return "\n".join(log), report
+    return "\n".join(lines), report
 
 # -----------------------------
-# Gradio UI
+# UI
 # -----------------------------
-with gr.Blocks(title="FieldSpace — Prime-Only Machine (RNS+CRT)") as demo:
-    gr.Markdown("# FieldSpace — Prime-Only Machine (RNS+CRT) — Exactness Proofs")
-    gr.Markdown("This demo proves exactness for a tiny CNN and a one-head attention numerator block using only residues mod $\\{2^K, p_i\\}$ and a single CRT.")
-    with gr.Tab("CNN (Conv + ReLU/Poly + FC)"):
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# FieldSpace, Prime Only Machine mod $2^K$, Exactness Proofs, single file")
+    gr.Markdown("CNN and 1 head attention, all int64, exact mod 2^K reconstruction.")
+
+    with gr.Tab("CNN Proof"):
+        with gr.Row():
+            seed   = gr.Number(value=0, label="Seed", precision=0)
+            xmax   = gr.Number(value=31, label="|X|max", precision=0)
+            wmax   = gr.Number(value=31, label="|W|max", precision=0)
+        act_sel   = gr.Radio(choices=["relu","poly"], value="relu", label="Activation")
+        shift_poly= gr.Slider(1, 12, value=7, step=1, label="SHIFT for poly x^2 >> SHIFT")
+        run_btn   = gr.Button("Run CNN Proof", variant="primary")
+        out_txt   = gr.Textbox(label="Console", lines=12)
+        out_json  = gr.JSON(label="JSON Report")
+        def _run_cnn(s, xM, wM, act, sh): return run_cnn_proof(int(s), int(xM), int(wM), act, int(sh))
+        run_btn.click(_run_cnn, [seed,xmax,wmax,act_sel,shift_poly], [out_txt,out_json], api_name="run_cnn_proof")
+
+    with gr.Tab("Attention Proof 1 head"):
+        with gr.Row():
+            seedA = gr.Number(value=0, label="Seed", precision=0)
+            XmaxA = gr.Number(value=7, label="X in [0, Xmax]", precision=0)
+            WmaxA = gr.Number(value=7, label="W in [0, Wmax]", precision=0)
         with gr.Row():
-            seed = gr.Number(value=0, precision=0, label="Seed")
-            xmax = gr.Number(value=31, precision=0, label="|X|max")
-            wmax = gr.Number(value=31, precision=0, label="|W|max")
-        activation = gr.Radio(choices=["relu","poly"], value="relu", label="Activation")
-        shift = gr.Slider(1,12,value=7,step=1,label="SHIFT for poly (x^2 >> SHIFT)")
-        run_btn = gr.Button("Run CNN Proof", variant="primary")
-        out_text = gr.Textbox(label="Console", lines=10)
-        out_json = gr.JSON(label="JSON Report")
-        run_btn.click(fn=lambda a,b,c,d,e: cnn_proof(int(a),int(b),int(c),d,int(e)),
-                      inputs=[seed,xmax,wmax,activation,shift],
-                      outputs=[out_text,out_json])
-    with gr.Tab("Attention (Exact Numerators)"):
-        aseed = gr.Number(value=0, precision=0, label="Seed")
-        runA = gr.Button("Run Attention Proof", variant="secondary")
-        AT = gr.Textbox(label="Console", lines=8)
-        AJ = gr.JSON(label="JSON Report")
-        runA.click(fn=lambda s: attn_numerators_proof(int(s)),
-                   inputs=[aseed], outputs=[AT,AJ])
-    gr.Markdown("**Use via API**")
-    gr.Code("""from gradio_client import Client
-client = Client('https://huggingface.co/spaces/jackal79/fieldspace-prime-only')
-txt, rep = client.predict('/run_cnn_proof', 0, 31, 31, 'relu', 7)""")
-
-# Named routes for programmatic calls
-demo.load(None, None, None)
-demo.add_named_endpoint("/run_cnn_proof", cnn_proof, inputs=[
-    gr.Number(precision=0), gr.Number(precision=0), gr.Number(precision=0),
-    gr.Textbox(), gr.Number(precision=0)
-], outputs=[gr.Textbox(), gr.JSON()])
-demo.add_named_endpoint("/run_attn_proof", attn_numerators_proof,
-                        inputs=[gr.Number(precision=0)],
-                        outputs=[gr.Textbox(), gr.JSON()])
+            BA = gr.Number(value=1, label="Batch B", precision=0)
+            TA = gr.Number(value=8, label="Seq T", precision=0)
+            CA = gr.Number(value=16, label="Width C", precision=0)
+            HA = gr.Number(value=1, label="Heads H divides C", precision=0)
+        runA   = gr.Button("Run Attention Proof", variant="primary")
+        outA_t = gr.Textbox(label="Console", lines=12)
+        outA_j = gr.JSON(label="JSON Report")
+        def _run_attn(s, xM, wM, B,T,C,H): return run_attn_proof(int(s), int(xM), int(wM), int(B), int(T), int(C), int(H))
+        runA.click(_run_attn, [seedA,XmaxA,WmaxA,BA,TA,CA,HA], [outA_t,outA_j], api_name="run_attn_proof")
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.launch()