#!/usr/bin/env python3 """Fits the lazy-eval probe (docs/NNUE.md's "Lazy-evaluation tiering plan"): a ridge regression on the trained net's own accumulator (fc1's input -- the ReLU'd, concatenated 512-wide per-perspective halves), targeting that SAME checkpoint's raw forward() output (pre-calibration -- see export_nnue_weights.py's CalibrateScoreMapping; NnueProbeScore applies the identical (CALIB_A, CALIB_B) correction itself at runtime, so the probe should target raw, not the calibrated cp value, to stay consistent with that). Must be re-run for every new checkpoint -- these weights are fit against one specific net's accumulator and are not portable across generations (same rule as material/score calibration). Writes the NPRB format NnueProbeLoadWeights (src/nnueeval.c) reads: magic: "NPRB" dim: uint32 (2*L1) mu: dim x float32 sigma: dim x float32 weight: dim x float32 bias: 1x float32 Usage: python3 fit_lazy_probe.py \ runs/sf_20261020_recal/nnue_ckpt/ckpt_epoch015.pt \ runs/sf_20261020_recal/nnue_probe.bin \ --train-tensor-dir runs/sf_20261020_recal/tensors_train \ --holdout-tensor-dir runs/sf_20261020_recal/tensors_holdout """ import argparse import struct from pathlib import Path import numpy as np import torch from nnue_model import NnueNet RIDGE_LAMBDA = 1.0 def _batched_accumulator_and_raw(model, tensor_dir, rows, batch_size=5000): """Returns (accum [N, 2*L1] float64, raw [N] float64) for the given row indices into tensor_dir's CSR arrays.""" d = Path(tensor_dir) white_idx = np.load(d / "white_idx.npy"); white_off = np.load(d / "white_off.npy") black_idx = np.load(d / "black_idx.npy"); black_off = np.load(d / "black_off.npy") stm = np.load(d / "stm.npy") accums, raws = [], [] with torch.no_grad(): for i in range(0, len(rows), batch_size): chunk = rows[i:i + batch_size] w_idx_list, w_off_list, b_idx_list, b_off_list = [], [], [], [] off_w, off_b = 0, 0 for r in chunk: wi = white_idx[white_off[r]:white_off[r + 1]] bi = black_idx[black_off[r]:black_off[r + 1]] w_idx_list.append(wi); w_off_list.append(off_w); off_w += len(wi) b_idx_list.append(bi); b_off_list.append(off_b); off_b += len(bi) w_idx = torch.from_numpy(np.concatenate(w_idx_list).astype(np.int64)) w_off = torch.tensor(w_off_list, dtype=torch.int64) b_idx = torch.from_numpy(np.concatenate(b_idx_list).astype(np.int64)) b_off = torch.tensor(b_off_list, dtype=torch.int64) stm_t = torch.from_numpy(stm[chunk].astype(np.int64)) acc = model.accumulator(w_idx, w_off, b_idx, b_off, stm_t) raw = model.forward(w_idx, w_off, b_idx, b_off, stm_t) accums.append(acc.numpy()) raws.append(raw.numpy()) return (np.concatenate(accums).astype(np.float64), np.concatenate(raws).astype(np.float64)) def main(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("checkpoint") ap.add_argument("outfile") ap.add_argument("--train-tensor-dir", required=True, help="fit the ridge regression against positions from here") ap.add_argument("--holdout-tensor-dir", required=True, help="report residual stats (not used for fitting) against here") ap.add_argument("--n-fit", type=int, default=80000) ap.add_argument("--n-report", type=int, default=20000) ap.add_argument("--seed", type=int, default=11) args = ap.parse_args() model = NnueNet() ck = torch.load(args.checkpoint, map_location="cpu", weights_only=True) model.load_state_dict(ck["model"]) model.eval() rng = np.random.default_rng(args.seed) n_train = len(np.load(Path(args.train_tensor_dir) / "stm.npy")) fit_rows = np.sort(rng.choice(n_train, size=min(args.n_fit, n_train), replace=False)) X, y = _batched_accumulator_and_raw(model, args.train_tensor_dir, fit_rows) mu = X.mean(axis=0) sigma = X.std(axis=0) sigma[sigma < 1e-6] = 1e-6 # guard against a dead/always-zero accumulator dim Xs = (X - mu) / sigma # Ridge regression, closed form: w = (X^T X + lambda*I)^-1 X^T y, # bias handled by centering y (mu/sigma already centers X). y_mean = y.mean() yc = y - y_mean dim = Xs.shape[1] A = Xs.T @ Xs + RIDGE_LAMBDA * np.eye(dim) w = np.linalg.solve(A, Xs.T @ yc) bias = y_mean # Xs has zero mean per-column, so w contributes 0 at the mean point fit_pred = Xs @ w + bias fit_resid = y - fit_pred print(f"Fit on {len(fit_rows)} positions ({args.train_tensor_dir}): " f"mean|resid|={np.mean(np.abs(fit_resid)):.3f}, " f"p98={np.percentile(np.abs(fit_resid), 98):.3f}") n_holdout = len(np.load(Path(args.holdout_tensor_dir) / "stm.npy")) report_rows = np.sort(rng.choice(n_holdout, size=min(args.n_report, n_holdout), replace=False)) Xh, yh = _batched_accumulator_and_raw(model, args.holdout_tensor_dir, report_rows) Xhs = (Xh - mu) / sigma holdout_pred = Xhs @ w + bias holdout_resid = yh - holdout_pred abs_resid = np.abs(holdout_resid) print(f"Held out on {len(report_rows)} positions ({args.holdout_tensor_dir}): " f"mean|resid|={np.mean(abs_resid):.3f}, " f"p50={np.percentile(abs_resid, 50):.3f}, " f"p98={np.percentile(abs_resid, 98):.3f}, " f"max={np.max(abs_resid):.3f}") print("NOTE: this is an OFFLINE, quiet-filtered-game-position sample -- " "NOT real search-node traffic. See the 2026-10-20 conversation on " "why that distinction matters before trusting these residual " "numbers as a safety bound for a real skip-gate.") with open(args.outfile, "wb") as f: f.write(b"NPRB") f.write(struct.pack("