Highly experimental

segmentation.py

@@ -150,10 +150,11 @@ def q2k_proxy_cfg(clamp_val: float, cfg: HPCConfig):
     """Cheap phase1+2 proxy for Q2_K label iteration (GGML seeds + WLS)."""
     from .q2k import codes_from_params_batch, dequantize_params_batch
     from .seeds import build_seeds
+    from .pipeline import _phase2_wls
 
     def _proxy(xg: np.ndarray, wg: np.ndarray | None) -> np.ndarray:
         seeds = build_seeds(xg, wg, cfg, clamp_val)
-        d, m = _wls(xg, wg, seeds, cfg)
+        d, m = _phase2_wls(xg, wg, seeds, cfg)
         sd = np.where(d > 1.0e-15, d, 1.0).astype(np.float32)
         sm = np.where(m > 1.0e-15, m, 1.0).astype(np.float32)
         ls = np.where(
@@ -172,57 +173,6 @@ def q2k_proxy_cfg(clamp_val: float, cfg: HPCConfig):
     return _proxy
 
 
-def _wls(xg, wg, seeds, cfg):
-    """Row-level WLS (same iteration as pipeline._phase2_wls, kept local to
-    avoid a circular import)."""
-    from .fp16 import round_fp16
-
-    n = xg.shape[0]
-    per_blk_clamp = np.zeros((n, 1), dtype=np.float32)
-    clipped = xg.astype(np.float32)
-    base_d = seeds["d"].astype(np.float64).copy()
-    base_m = seeds["dmin"].astype(np.float64).copy()
-    ls = seeds["ls"].copy()
-    lm = seeds["lm"].copy()
-    scales = seeds["scales"].astype(np.float64)
-    mins = seeds["mins"].astype(np.float64)
-    ww = np.ones_like(xg, dtype=np.float64) if wg is None else wg.astype(np.float64)
-
-    from .q2k import codes_from_params_batch
-
-    for _ in range(5):
-        codes = codes_from_params_batch(clipped, base_d, base_m, ls, lm).astype(np.float64)
-        A = np.repeat(ls.astype(np.float64), 16, axis=1) * codes
-        B = np.repeat(lm.astype(np.float64), 16, axis=1)
-        Saa = np.sum(ww * A * A, axis=1)
-        Sab = np.sum(ww * A * B, axis=1)
-        Sbb = np.sum(ww * B * B, axis=1)
-        Sxa = np.sum(ww * clipped.astype(np.float64) * A, axis=1)
-        Sxb = np.sum(ww * clipped.astype(np.float64) * B, axis=1)
-        det = Saa * Sbb - Sab * Sab
-        valid = det > 1.0e-30
-        safe_det = np.where(valid, det, 1.0)
-        d_new = (Sbb * Sxa - Sab * Sxb) / safe_det
-        m_new = (Sab * Sxa - Saa * Sxb) / safe_det
-        upd_d = valid & (d_new > 0.0) & np.isfinite(d_new)
-        upd_m = valid & (m_new > 0.0) & np.isfinite(m_new)
-        base_d = np.where(upd_d, round_fp16(d_new.astype(np.float32)).astype(np.float64), base_d)
-        base_m = np.where(upd_m, round_fp16(m_new.astype(np.float32)).astype(np.float64), base_m)
-        safe_d = np.where(base_d > 1.0e-15, base_d, 1.0)
-        safe_m = np.where(base_m > 1.0e-15, base_m, 1.0)
-        ls = np.where(
-            base_d[:, None] > 1.0e-15,
-            np.clip(nearest_int(scales / safe_d[:, None]), 0, 15),
-            0,
-        ).astype(np.uint8)
-        lm = np.where(
-            base_m[:, None] > 1.0e-15,
-            np.clip(nearest_int(mins / safe_m[:, None]), 0, 15),
-            0,
-        ).astype(np.uint8)
-    return base_d, base_m
-
-
 def q3k_proxy(xg: np.ndarray, wg: np.ndarray | None) -> np.ndarray:
     """Cheap per-16 sub-block 3-bit ladder proxy for Q3_K label iteration."""
     from .seeds import make_qkx2_batch