"""Geometry-preserving weight analogs for Q2_K bilinear factorization. At 2 bits per weight (4 levels ``q in {0, 1, 2, 3}`` and 4-bit sub-block scales ``Ls, Lm in {0..15}``), raw weight blocks suffer from two structural bottlenecks: 1. **Intra-sub-block outlier collapse**: a single heavy-tailed weight stretches the sub-block span ``xmax - xmin``, forcing 13-15 of the 16 weights into one or two quantization bins. 2. **Inter-sub-block scale crushing**: a single high-variance sub-block inflates the shared superblock scale ``d``, collapsing ``Ls_j`` of the remaining 15 sub-blocks toward small integers with severe rounding loss. Rather than factoring raw weights alone, V2 constructs a bank of **geometry-preserving weight analogs** ``x_tilde`` of the eligible weights: * :func:`subblock_isomoment_analog` maps each 16-weight sub-block through an importance-weighted saturating transport blended with the optimal 4-level quantile spacing, then orthogonally projects back onto the exact weighted mean ``mu_j`` and weighted variance ``sigma_j^2`` of the original sub-block. * :func:`vesica_gram_analog` operates in the D6 / Vesica fold basis ``(x_p + x_{p+128}, x_p - x_{p+128})``, regularizing outliers in the whitened eigen-basis and coloring back so the exact 2x2 weighted Gram matrix and means of the two half-blocks are preserved. * :func:`subscale_lattice_analog` tempers cross-sub-block dispersion toward the 4-bit ``{0..15}`` scale lattice while preserving every sub-block's unit direction, mean, and total block weighted energy. The discrete Q2_K lattice coordinates ``(Ls, Lm, q)`` discovered on each analog are then projected back onto the **true weights** ``x`` via coupled bilinear least-squares, guaranteeing that the final emitted block is strictly optimal under the true objective. """ from __future__ import annotations import numpy as np from .q2k import NS, QK_K _Q2_IDEAL_QUANTILES = np.linspace(-1.35, 1.35, 16, dtype=np.float64) def subblock_isomoment_analog( x: np.ndarray, w: np.ndarray | None = None, tau: float = 1.65, beta: float = 0.25, quantile_blend: float = 0.30, enforce_nonpositive_min: bool = False, ) -> np.ndarray: """Create a sub-block iso-moment analog preserving weighted mean & variance. Invariants per 16-weight sub-block ``j`` (when ``enforce_nonpositive_min`` is ``False``): * ``sum_k w_{j,k} x_tilde_{j,k} == sum_k w_{j,k} x_{j,k}`` * ``sum_k w_{j,k} (x_tilde_{j,k} - mu_j)^2 == sum_k w_{j,k} (x_{j,k} - mu_j)^2`` * strict rank-order preservation with respect to ``x_{j,k}`` (when ``beta == 0``) """ xb = np.asarray(x, dtype=np.float64).reshape(NS, 16) wb = ( np.ones((NS, 16), dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64).reshape(NS, 16) ) wsum = np.maximum(np.sum(wb, axis=1, keepdims=True), 1.0e-30) mu = np.sum(wb * xb, axis=1, keepdims=True) / wsum centered = xb - mu var = np.sum(wb * centered * centered, axis=1, keepdims=True) / wsum sigma = np.sqrt(np.maximum(var, 1.0e-30)) active = (sigma[:, 0] >= 1.0e-15) z = centered / sigma w_mean = wsum / 16.0 w_rel = np.power(np.maximum(wb / w_mean, 1.0e-12), beta) if beta != 0.0 else 1.0 z_comp = tau * np.tanh((z * w_rel) / max(tau, 1.0e-6)) if quantile_blend > 0.0: order = np.argsort(z, axis=1, kind="stable") z_quant = np.empty_like(z) rows = np.arange(NS)[:, None] z_quant[rows, order] = _Q2_IDEAL_QUANTILES[None, :] z_analog = (1.0 - quantile_blend) * z_comp + quantile_blend * z_quant else: z_analog = z_comp mu_a = np.sum(wb * z_analog, axis=1, keepdims=True) / wsum z_analog = z_analog - mu_a var_a = np.sum(wb * z_analog * z_analog, axis=1, keepdims=True) / wsum scale = np.where(var_a > 1.0e-30, np.sqrt(var / np.maximum(var_a, 1.0e-30)), 0.0) out = np.where(active[:, None], mu + z_analog * scale, xb) if enforce_nonpositive_min: orig_min = np.min(xb, axis=1) new_min = np.min(out, axis=1) shift = np.where((orig_min <= 0.0) & (new_min > 0.0), new_min, 0.0) out = out - shift[:, None] return out.reshape(QK_K).astype(np.float32) def vesica_gram_analog( x: np.ndarray, w: np.ndarray | None = None, tau: float = 1.9, ) -> np.ndarray: """Create a D6 / Vesica fold analog preserving the 2x2 half-block Gram matrix. Decomposes ``x`` into symmetric ``v = 0.5*(x[:128] + x[128:])`` and antisymmetric ``u = 0.5*(x[:128] - x[128:])`` modes under pair weights ``wp = 0.5*(w[:128] + w[128:])``. Whitens ``(v, u)``, applies smooth hyperbolic saturation in the principal axes, and colors back with the exact 2x2 covariance square-root so ``Gram(v_tilde, u_tilde) == Gram(v, u)``. """ x_f64 = np.asarray(x, dtype=np.float64).reshape(QK_K) w_f64 = ( np.ones(QK_K, dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64).reshape(QK_K) ) half = QK_K // 2 x0, x1 = x_f64[:half], x_f64[half:] wp = np.maximum(0.5 * (w_f64[:half] + w_f64[half:]), 1.0e-30) wsum = float(np.sum(wp)) v = 0.5 * (x0 + x1) u = 0.5 * (x0 - x1) mv = float(np.sum(wp * v) / wsum) mu = float(np.sum(wp * u) / wsum) Y = np.vstack([v - mv, u - mu]) # (2, 128) G = (Y * wp[None, :]) @ Y.T / wsum # 2x2 weighted covariance evals, evecs = np.linalg.eigh(G) if float(np.max(evals)) < 1.0e-28: return x_f64.astype(np.float32) pos = evals > 1.0e-28 inv_sqrt = np.where(pos, 1.0 / np.sqrt(np.maximum(evals, 1.0e-30)), 0.0) W_white = (evecs * inv_sqrt[None, :]) @ evecs.T Z = W_white @ Y # whitened (2, 128) Z_comp = tau * np.tanh(Z / max(tau, 1.0e-6)) Z_mean = np.sum(Z_comp * wp[None, :], axis=1, keepdims=True) / wsum Z_comp = Z_comp - Z_mean G_z = (Z_comp * wp[None, :]) @ Z_comp.T / wsum ez, Vz = np.linalg.eigh(G_z) pos_z = ez > 1.0e-28 if not np.all(pos_z): return x_f64.astype(np.float32) inv_sqrt_z = np.where(pos_z, 1.0 / np.sqrt(np.maximum(ez, 1.0e-30)), 0.0) Z_orth = (Vz * inv_sqrt_z[None, :]) @ Vz.T @ Z_comp sqrt_G = (evecs * np.sqrt(np.maximum(evals, 0.0))[None, :]) @ evecs.T Y_tilde = sqrt_G @ Z_orth v_new = mv + Y_tilde[0] u_new = mu + Y_tilde[1] return np.concatenate([v_new + u_new, v_new - u_new]).astype(np.float32) def subscale_lattice_analog( x: np.ndarray, w: np.ndarray | None = None, gamma: float = 0.65, ) -> np.ndarray: """Create a cross-sub-block analog whose sub-block dispersions fit 4-bit scales. Preserves each sub-block's mean ``mu_j`` and unit direction ``(x_j - mu_j) / sigma_j``, while tempering extreme sub-block standard deviations ``sigma_j`` toward the superblock root-mean-square dispersion and rescaling to preserve total centered weighted energy. """ xb = np.asarray(x, dtype=np.float64).reshape(NS, 16) wb = ( np.ones((NS, 16), dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64).reshape(NS, 16) ) wsum = np.maximum(np.sum(wb, axis=1, keepdims=True), 1.0e-30) mu = np.sum(wb * xb, axis=1, keepdims=True) / wsum centered = xb - mu energy_j = np.sum(wb * centered * centered, axis=1, keepdims=True) total_energy = float(np.sum(energy_j)) if total_energy <= 1.0e-30: return xb.reshape(QK_K).astype(np.float32) sigma = np.sqrt(np.maximum(energy_j / wsum, 1.0e-30)) sigma_ref = float(np.sqrt(total_energy / max(float(np.sum(wsum)), 1.0e-30))) target_sigma = np.power(np.maximum(sigma, 1.0e-30), gamma) * ( sigma_ref ** (1.0 - gamma) ) scaled = np.where(sigma > 1.0e-15, centered * (target_sigma / sigma), 0.0) new_energy = float(np.sum(wb * scaled * scaled)) if new_energy > 1.0e-30: scaled = scaled * np.sqrt(total_energy / new_energy) out = mu + scaled return out.reshape(QK_K).astype(np.float32) def build_block_analogs( x: np.ndarray, w: np.ndarray | None = None, ) -> list[np.ndarray]: """Return the portfolio of geometry-preserving weight analogs for a block.""" a1 = subblock_isomoment_analog( x, w, tau=1.60, beta=0.25, quantile_blend=0.30, enforce_nonpositive_min=True ) a2 = subblock_isomoment_analog( x, w, tau=2.15, beta=0.15, quantile_blend=0.15, enforce_nonpositive_min=True ) a3 = vesica_gram_analog(x, w, tau=1.85) a4 = subscale_lattice_analog(a1, w, gamma=0.70) return [a1, a2, a3, a4]