"""Ternary weight quantisers (spec §4): Q-A absmean, Q-B TWN, Q-C learnable (LSQ-style) scale.

Weights have shape [..., in]; groups of `group_size` run along the input (last) dimension, matching
llama.cpp's block layout. Stacked expert tensors [E, out, in] work unchanged.

The forward weight is s * q with q in {-1, 0, +1}. It is built by a custom autograd function rather
than `w + (w_hat - w).detach()`: in fp32 that expression is not always bit-equal to w_hat, and at
7.7B weights the difference flips ~1e5 bf16 values. Gradients to the latent are the identity (plain
STE) or zero where |w|/s > clip_threshold (clipped STE).

Scales are canonicalised to bf16 precision (DECISIONS.md P-03): the trainer's matmul sees
bf16(s * q) = bf16(s) * q, and bf16(s) is exactly representable in the fp16 export format, so the
fake-quant, packed and GGUF students share one set of weights by construction.

D-51 adds three N-bit expert formats, each trained on exactly the grid of a llama.cpp GGUF type so the export is
lossless (section "N-bit grids" below): int2 = Q2_0, int4 = Q4_0, int3 = Q3_K (its down projection Q4_0).
D-57 adds llama.cpp's non-linear 4-bit grid: iq4nl = IQ4_NL, iq4xs = IQ4_XS (its down projection IQ4_NL), and
int3nl = Q3_K with its down projection IQ4_NL.
"""
from __future__ import annotations

import math
from dataclasses import dataclass

import torch

METHODS = ("absmean", "twn", "lsq", "int2", "int3", "int4",  # int*: D-51, see "N-bit grids"
           "int3nl", "iq4nl", "iq4xs", "int2nl",  # D-57
           "codebook")  # D-58: llama.cpp's codebook formats, frozen codes (`quant.codebook`)
STES = ("plain", "clipped")
TWN_DELTA = 0.7


@dataclass(frozen=True)
class QuantSpec:
    method: str = "absmean"  # absmean (Q-A) | twn (Q-B) | lsq (Q-C)
    group_size: int = 128
    ste: str = "plain"  # plain | clipped
    clip_threshold: float = 1.5
    eps: float = 1e-8

    def __post_init__(self) -> None:
        if self.method not in METHODS:
            raise ValueError(f"method {self.method!r} not in {METHODS}")
        if self.ste not in STES:
            raise ValueError(f"ste {self.ste!r} not in {STES}")
        if self.group_size <= 0:
            raise ValueError("group_size must be positive")
        if self.method in GRIDS:
            if self.group_size != GRIDS[self.method].group:
                raise ValueError(f"{self.method} is GGUF {GRIDS[self.method].gguf}: group size must be "
                                 f"{GRIDS[self.method].group}, not {self.group_size}")
            if self.ste != "plain":
                raise ValueError(f"{self.method}: only the plain STE is built")


def _grouped(w: torch.Tensor, group: int) -> torch.Tensor:
    if w.shape[-1] % group:
        raise ValueError(f"input dim {w.shape[-1]} not divisible by group size {group}")
    return w.unflatten(-1, (-1, group))


def canonical_scale(s: torch.Tensor) -> torch.Tensor:
    """Round scales to bf16 precision (nearest, ties to even), returned as fp32 (exactly representable in fp16).

    Done on the fp32 bits: torch.compile elides a `.to(bfloat16).to(float32)` round trip inside a fused kernel,
    which left every scale unrounded (T2.3). Scales are positive and finite, so the bit trick is exact."""
    b = s.to(torch.float32).view(torch.int32)
    b = (b + 0x7FFF + ((b >> 16) & 1)) & -0x10000
    return b.view(torch.float32)


def init_scale(w: torch.Tensor, spec: QuantSpec, importance: torch.Tensor | None = None) -> torch.Tensor:
    """Per-group scale [..., in/g] in fp32, canonicalised. Q-C initialises from Q-A. The N-bit methods return
    their (signed) per-group scale, `nbit_scale`.

    Group sums accumulate in fp64, as the reference does, so the scale does not depend on reduction order: under
    torch.compile fp32 sums moved scales across bf16 rounding boundaries (0.085% of Marco's codes, T2.3)."""
    if spec.method in GRIDS:
        return nbit_scale(w, spec, importance)
    wg = _grouped(w.detach().double(), spec.group_size).abs()
    if spec.method == "twn":
        delta = TWN_DELTA * wg.mean(-1, keepdim=True)
        above = wg > delta
        s = (wg * above).sum(-1) / above.sum(-1).clamp_min(1)
    else:
        s = wg.mean(-1)
    return canonical_scale(s.float().clamp_min(spec.eps))


def ternary_codes(w: torch.Tensor, s: torch.Tensor, spec: QuantSpec) -> torch.Tensor:
    """int8 codes in {-1, 0, 1}, shape of w.

    Round-half-even of w/s clamped to [-1, 1] is sign(w)·[|w| > s/2]. With s bf16-exact, s/2 is exact, so the
    comparison is exact where an fp32 division is not (compiled kernels may divide differently).
    The N-bit methods return their grid codes (`nbit_codes`)."""
    if spec.method in GRIDS:
        return nbit_codes(w, s, spec)
    wg = _grouped(w.detach().float(), spec.group_size)
    if spec.method == "twn":
        a = wg.abs()
        delta = TWN_DELTA * a.double().mean(-1, keepdim=True)  # fp64, as in init_scale
        q = torch.sign(wg) * (a > delta)
    else:
        q = torch.sign(wg) * (wg.abs() > 0.5 * s.unsqueeze(-1))
    return q.flatten(-2).to(torch.int8)


class _TernarySTE(torch.autograd.Function):
    """Forward s*q exactly; backward identity (or clipped) to w, LSQ gradient to a learnable s."""

    @staticmethod
    def forward(ctx, w, s, spec: QuantSpec):
        codes = ternary_codes(w, s, spec)
        s_c = canonical_scale(s.detach().float())
        w_hat = (_grouped(codes.to(torch.float32), spec.group_size) * s_c.unsqueeze(-1)).flatten(-2)
        ctx.spec = spec
        ctx.save_for_backward(w, s_c, codes)  # int8: a quarter of fp32 codes (Marco: 0.27 vs 1.1 GB per layer)
        return w_hat.to(w.dtype)

    @staticmethod
    def backward(ctx, g):
        w, s_c, codes = ctx.saved_tensors
        spec: QuantSpec = ctx.spec
        gw = g
        ratio = None
        if spec.ste == "clipped" or (spec.method == "lsq" and ctx.needs_input_grad[1]):
            ratio = _grouped(w.float(), spec.group_size) / s_c.unsqueeze(-1)
        if spec.ste == "clipped":
            gw = g * (ratio.abs() <= spec.clip_threshold).flatten(-2).to(g.dtype)
        gs = None
        if spec.method == "lsq" and ctx.needs_input_grad[1]:
            # LSQ: d(s*q)/ds = q - w/s inside the rounding range, q outside; grad scale 1/sqrt(g*Qp), Qp = 1
            qg = _grouped(codes.float(), spec.group_size)
            inside = ratio.abs() <= 1.5
            dws = torch.where(inside, qg - ratio, qg)
            gs = (_grouped(g.float(), spec.group_size) * dws).sum(-1) / math.sqrt(spec.group_size)
        return gw, gs, None


def fake_quant(w: torch.Tensor, spec: QuantSpec, scale: torch.Tensor | None = None,
               importance: torch.Tensor | None = None) -> torch.Tensor:
    """Ternary (or D-51 N-bit) forward weight with STE. `scale` is required (a trained parameter) for lsq.
    D-57 `importance` [..., in] (w without its output dim) weights the N-bit scale search (`_error_weight`)."""
    if spec.method in GRIDS:
        return _NBitSTE.apply(w, spec, importance)
    if spec.method == "lsq":
        if scale is None:
            raise ValueError("lsq needs a learnable scale (initialise with init_scale)")
        s = scale
    else:
        s = init_scale(w, spec)
    return _TernarySTE.apply(w, s, spec)


def scale_for(w: torch.Tensor, spec: QuantSpec, scale: torch.Tensor | None = None) -> torch.Tensor:
    """The canonical scale actually used in the forward pass."""
    return canonical_scale(scale.detach().float()) if spec.method == "lsq" else init_scale(w, spec)


# ---------------------------------------------------------------- packing


def pack_2bit(codes: torch.Tensor) -> torch.Tensor:
    """{-1,0,1} int8 [..., n] -> uint8 [..., n/4]; element i in bits 2*(i%4) of byte i//4, stored as q+1."""
    if codes.shape[-1] % 4:
        raise ValueError("last dim must be divisible by 4")
    u = (codes.to(torch.int16) + 1).to(torch.uint8).unflatten(-1, (-1, 4))
    return u[..., 0] | (u[..., 1] << 2) | (u[..., 2] << 4) | (u[..., 3] << 6)


def unpack_2bit(packed: torch.Tensor) -> torch.Tensor:
    parts = [(packed >> (2 * i)) & 3 for i in range(4)]
    return (torch.stack(parts, -1).flatten(-2).to(torch.int16) - 1).to(torch.int8)


_POW3 = (1, 3, 9, 27, 81)


def pack_trits(codes: torch.Tensor) -> torch.Tensor:
    """{-1,0,1} int8 [..., n] -> uint8 [..., ceil(n/5)]; byte = sum_i (q_i+1) * 3^i (5 trits/byte, 1.6 bits)."""
    n = codes.shape[-1]
    pad = (-n) % 5
    t = (codes.to(torch.int16) + 1)
    if pad:
        t = torch.nn.functional.pad(t, (0, pad), value=1)  # pad with code 0
    t = t.unflatten(-1, (-1, 5))
    pw = torch.tensor(_POW3, dtype=torch.int16, device=codes.device)
    return (t * pw).sum(-1).to(torch.uint8)


def unpack_trits(packed: torch.Tensor, n: int) -> torch.Tensor:
    v = packed.to(torch.int16)
    out = []
    for _ in range(5):
        out.append(v % 3)
        v = v // 3
    t = torch.stack(out, -1).flatten(-2)[..., :n]
    return (t - 1).to(torch.int8)


def export_packed(w: torch.Tensor, spec: QuantSpec, scale: torch.Tensor | None = None,
                  importance: torch.Tensor | None = None) -> dict:
    """Packed student tensor: 2-bit codes and fp16 scales (exactly the forward's bf16 scales). N-bit methods:
    `export_packed_nbit` (with the forward's `importance`, so the codes are the trained ones)."""
    if spec.method in GRIDS:
        return export_packed_nbit(w, spec, importance)
    s = scale_for(w, spec, scale)
    return {"codes": pack_2bit(ternary_codes(w, s, spec)), "scales": s.to(torch.float16),
            "group_size": spec.group_size, "shape": tuple(w.shape)}


def dequant_packed(p: dict, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
    if p.get("method") in GRIDS:
        return dequant_packed_nbit(p, dtype)
    q = unpack_2bit(p["codes"]).to(torch.float32)
    w = (_grouped(q, p["group_size"]) * p["scales"].to(torch.float32).unsqueeze(-1)).flatten(-2)
    return w.to(dtype)


# ---------------------------------------------------------------- N-bit grids (D-51)
#
# Each N-bit format is the value grid of one llama.cpp GGUF type, so the GGUF export is a repack, not a quantisation:
#   int2 = Q2_0: blocks of 64, fp16 d, value q·d with q in {-1, 0, 1, 2}.
#   int4 = Q4_0: blocks of 32, fp16 d, value q·d with q in -8..7.
#   int3 = Q3_K: super-blocks of 256 with an fp16 d; 16 sub-blocks of 16, each with a 6-bit scale m in -32..31;
#          value d·m·q with q in -4..3 (ggml's dequantize_row_q3_K). Q3_K needs rows of a multiple of 256, and
#          Marco's expert down projection has rows of 384 (moe_intermediate_size), so int3 trains its down
#          projection on int4's grid and exports it as Q4_0 (`proj_spec`): int3 is a mixed 3/4-bit model.
#
# Scales are signed (both GGUF types store a signed d), so each group puts its extreme weight E (the largest |w|,
# the positive one on a tie) on the long end of its asymmetric grid: s = E / ext with ext = +2 (int2), -8 (int4),
# -4 (int3's sub-blocks), which for c = 1 is llama.cpp's own Q4_0 rule. A fixed search then tries s = c·E/ext for
# the few c in `cands` and keeps the one with the least squared error over the group (the first on a tie); int3
# instead tries m = round(|t|/d) + k for the offsets k in `cands` (below). On Marco's experts (scratch measurement
# on 6-9 expert tensors of layers 2, 14, 26, 2026-10-22) the search cuts the round-to-nearest error against c = 1
# (k = 0) alone by 19% (int2), 6% (int4) and 12% (int3), to within 1-4% of an oracle search (151 values of c; every
# m in 1..31).
#
# bf16 exactness (P-03): the trainer's matmul sees bf16(s·q), so s·q must be exact in bf16 (8 significant bits)
# for every code. int2's codes have one significant bit, so its scales keep bf16's 8; int4's need up to three
# (5, 7), so its scales keep 5. int3's d·m·q needs d·m·3 in 8 bits, and m alone takes up to 5, so d is a power of
# two: the smallest at or above max|t|/31 over the super-block (t = the sub-block's E/ext), and
# m = round(|t|/d) + k clamped to 1..31 (offsets, not fractions: a small sub-block's m is a few units, where
# fractions of it round to the same m). Every |scale| (int3: d) is at least 2^-14, fp16's smallest normal, so all are
# exact in fp16 and the GGUF dequantises to the trained weights bit for bit (`export.gguf.check_round_trip`).
#
# Determinism (T2.3): group extremes are exact; scale arithmetic runs in fp64 and is rounded on the fp32 bits; codes
# are round-half-even of w/s made exact by checking an fp32 guess against the half-way points (q ± 0.5)·s, which
# are exact (≤ 13 significant bits); the search sums squared errors in fp64. Compiled and eager runs agree
# (tests/test_nbit.py). The gradient is the plain STE.


@dataclass(frozen=True)
class Grid:
    gguf: str
    lo: int
    hi: int
    group: int
    ext: int        # the grid end a group's extreme weight lands on (its sign orients the scale)
    sig: int        # significant bits kept in a scale (int3: unused, d is a power of two)
    cands: tuple    # the search's fractions c of E/ext (int3, iq4xs: integer offsets k to m)
    values: tuple | None = None  # D-57 non-linear grids: the value of each code lo..hi (None: the code itself)


# D-57: IQ4_NL's 16 levels (ggml's kvalues_iq4nl), shared by IQ4_XS. Codes are indices 0..15 into this table.
IQ4NL_VALUES = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113)
GRIDS = {"int2": Grid("Q2_0", -1, 2, 64, 2, 8, (1.0, 0.9, 0.8, 0.7, 0.6)),
         "int4": Grid("Q4_0", -8, 7, 32, -8, 5, (1.0, 15 / 16, 7 / 8)),
         "int3": Grid("Q3_K", -4, 3, 16, -4, 0, (-3, -2, -1, 0, 1)),
         "iq4nl": Grid("IQ4_NL", 0, 15, 32, -127, 11, (1.0, 0.96, 0.92, 0.88, 0.84, 0.8, 0.76), IQ4NL_VALUES),
         "iq4xs": Grid("IQ4_XS", 0, 15, 32, -127, 11, (-1, 0, 1), IQ4NL_VALUES)}
GRIDS["int3nl"] = GRIDS["int3"]  # D-57: Q3_K gate/up; only its down projection differs (`model.student.proj_spec`)
GRIDS["int2nl"] = GRIDS["int2"]  # D-57: Q2_0 gate/up with an IQ4_NL down projection (`model.student.proj_spec`)
Q3K_METHODS = ("int3", "int3nl")
INT2_METHODS = ("int2", "int2nl")
Q3K_SUPER, Q3K_M_MAX = 256, 31
IQ4XS_SUPER, IQ4XS_M_MAX = 256, 31
SCALE_MIN = 2.0 ** -14  # fp16's smallest normal
_SLICE = 1 << 24  # eager calls on stacked tensors (export, eval) run in slices of about this many weights


def _sliced(fn, w: torch.Tensor, s: torch.Tensor | None = None):
    """fn(w) or fn(w, s) on slices of dim 0 when eager: the search's temporaries are several times its input, and a
    Marco layer's gate_up is 182M weights. Groups never cross dim 0, so the result is the same."""
    call = (lambda a, b: fn(a)) if s is None else fn
    if torch.compiler.is_compiling() or w.dim() < 2 or w.numel() <= _SLICE:
        return call(w, s)
    n = max(1, _SLICE // (w.numel() // w.shape[0]))
    parts = [call(w[i:i + n], None if s is None else s[i:i + n]) for i in range(0, w.shape[0], n)]
    if isinstance(parts[0], tuple):
        return tuple(torch.cat(p) for p in zip(*parts))
    return torch.cat(parts)


def _round_sig(x: torch.Tensor, sig: int) -> torch.Tensor:
    """|x| rounded to `sig` significant bits (nearest, ties to even) on its fp32 bits, at least 2^-14; sign kept
    (-0 counts as positive). For positive x ≥ 2^-14 and sig = 8 this is `canonical_scale`."""
    a = x.float().abs().clamp_min(SCALE_MIN)
    drop = 24 - sig
    b = a.view(torch.int32)
    b = (b + ((1 << (drop - 1)) - 1) + ((b >> drop) & 1)) & -(1 << drop)
    a = b.view(torch.float32)
    return torch.where(x < 0, -a, a)


def _extreme(wg: torch.Tensor) -> torch.Tensor:
    """Each group's largest-magnitude weight, signed; the positive one on a tie."""
    mx, mn = wg.amax(-1), wg.amin(-1)
    return torch.where(mx >= -mn, mx, mn)


def _round_grid(x: torch.Tensor, a: torch.Tensor, lo: int, hi: int) -> torch.Tensor:
    """round-half-even(x / a) clamped to [lo, hi], exactly, for a > 0 with few significant bits: the fp32 quotient's
    guess is off by at most one, and comparing x with the exact half-way points (q ± 0.5)·a corrects it."""
    q = torch.round(x / a).clamp(lo - 1, hi + 1)
    up, dn = (q + 0.5) * a, (q - 0.5) * a
    odd = torch.remainder(q, 2) == 1  # at a tie the even neighbour wins: q ± 1 is even when q is odd
    q = q + ((x > up) | ((x == up) & odd)).to(x.dtype) - ((x < dn) | ((x == dn) & odd)).to(x.dtype)
    return q.clamp(lo, hi)


def _nearest_index(x: torch.Tensor, values: tuple) -> torch.Tensor:
    """D-57: index of the nearest table value to x (the lower one at a midpoint), as float codes 0..len-1."""
    v = torch.tensor(values, dtype=torch.float32, device=x.device)
    mids = (v[1:] + v[:-1]) / 2  # half-integers, exact
    return torch.bucketize(x.float(), mids).to(x.dtype)  # = (x > mids).sum(-1), by binary search


def _code_values(q: torch.Tensor, g: Grid) -> torch.Tensor:
    """The grid value of each code: the code itself, or (D-57) its entry in the table."""
    if g.values is None:
        return q
    return torch.tensor(g.values, dtype=torch.float32, device=q.device)[q.long()]


def _grid_codes(wg: torch.Tensor, s: torch.Tensor, g: Grid) -> torch.Tensor:
    """Codes of grouped weights [..., n, g] against signed scales [..., n]: round(w / s) on the grid (D-57: the
    nearest table value; the quotient is not exact there, which only moves weights within a hair of a midpoint)."""
    s = s.unsqueeze(-1)
    x = torch.where(s < 0, -wg, wg)
    if g.values is not None:
        return _nearest_index(x / s.abs(), g.values)
    return _round_grid(x, s.abs(), g.lo, g.hi)


def _search(wg: torch.Tensor, g: Grid, cands: list[torch.Tensor], unit: torch.Tensor | None = None,
            wt: torch.Tensor | None = None) -> torch.Tensor:
    """The candidate (a scale, or int3's m with scale unit·m) whose codes leave the least squared error per group
    (D-57: weighted by `wt`, broadcastable to wg, when given)."""
    best = err = None
    for c in cands:
        s = c if unit is None else unit * c
        q = _grid_codes(wg, s, g)
        d2 = (wg - _code_values(q, g) * s.unsqueeze(-1)).double().square()
        e = (d2 if wt is None else d2 * wt).sum(-1)
        if best is None:
            best, err = c, e
        else:
            take = e < err
            best, err = torch.where(take, c, best), torch.where(take, e, err)
    return best


def _error_weight(wg: torch.Tensor, imp: torch.Tensor | None, g: Grid) -> torch.Tensor | None:
    """D-57: llama.cpp's imatrix weighting of a group's squared errors, imp_j · sqrt(σ² + w_j²) with σ² = 2·mean(w²)
    over the group (its quantize_row_iq4_nl_impl). `imp` [..., in] is the mean squared input activation of each
    input column; w is [..., out, in], so imp broadcasts over the output rows. None without an importance."""
    if imp is None or g.values is None:  # only the non-linear IQ4 grids: Q2_0/Q4_0 keep their trained rounding
        return None
    ig = _grouped(imp.float().unsqueeze(-2), g.group)
    sigma2 = 2 * wg.square().mean(-1, keepdim=True)
    return (ig * (sigma2 + wg.square()).sqrt()).double()


def _flat_scale(w: torch.Tensor, g: Grid, imp: torch.Tensor | None = None) -> torch.Tensor:
    wg = _grouped(w, g.group)
    e = _extreme(wg).double()
    wt = _error_weight(wg, imp, g)
    s = _search(wg, g, [_round_sig(e * (c / g.ext), g.sig) for c in g.cands], wt=wt)
    if g.values is None:
        return s
    # D-57 non-linear grids: one (weighted) least-squares refit of the scale to the chosen codes, kept where it lowers
    # the error with its own nearest codes, as llama.cpp's IQ4 quantiser does (sumqx / sumq2)
    v = _code_values(_grid_codes(wg, s, g), g).double()
    ww = torch.ones_like(v) if wt is None else wt
    num, den = (ww * wg.double() * v).sum(-1), (ww * v * v).sum(-1)
    ls = torch.where(den > 0, num / den.clamp_min(1e-30), s.double())
    return _search(wg, g, [s, _round_sig(ls, g.sig)], wt=wt)


def _q3k_factors(w: torch.Tensor, g: Grid) -> tuple[torch.Tensor, torch.Tensor]:
    """Q3_K's scales of w [..., C]: d [..., C/256] (fp32 powers of two) and m [..., C/16] (signed integers in fp32,
    1 ≤ |m| ≤ 31)."""
    if w.shape[-1] % Q3K_SUPER:
        raise ValueError(f"int3 (Q3_K) needs rows of a multiple of {Q3K_SUPER}, not {w.shape[-1]}")
    wg = _grouped(w, g.group)
    t = _extreme(wg).double() / g.ext  # exact: ext is a power of two
    x = (_grouped(t.abs(), Q3K_SUPER // g.group).amax(-1) / Q3K_M_MAX).float().clamp_min(SCALE_MIN)
    d = ((x.view(torch.int32) + 0x7FFFFF) & 0x7F800000).view(torch.float32)  # the power of two at or above x
    unit = d.unsqueeze(-1).expand(*d.shape, Q3K_SUPER // g.group).flatten(-2)  # d per sub-block
    sign = torch.where(t < 0, -1.0, 1.0).float()
    m0 = (t.abs() / unit.double()).round()  # exact quotient (d is a power of two), half to even
    cands = [sign * (m0 + k).clamp(1, Q3K_M_MAX).float() for k in g.cands]
    return d, _search(wg, g, cands, unit)


def _q3k_scale(d: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
    return d.unsqueeze(-1).expand(*d.shape, Q3K_SUPER // GRIDS["int3"].group).flatten(-2) * m


def _iq4xs_factors(w: torch.Tensor, g: Grid, imp: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
    """D-57 IQ4_XS's scales of w [..., C]: d [..., C/256] (fp16-exact fp32) and m [..., C/32] (signed integers in
    fp32, 1 ≤ |m| ≤ 31; GGUF stores m + 32 in 6 bits). Each sub-block's own IQ4_NL search sets its target t; d puts
    the largest |t| of the super-block on m = 31, and m = round(|t|/d) + k for k in `cands` keeps the least error.
    d·m has at most 16 significant bits, so d·m·value is exact in fp32, as ggml computes it."""
    if w.shape[-1] % IQ4XS_SUPER:
        raise ValueError(f"iq4xs (IQ4_XS) needs rows of a multiple of {IQ4XS_SUPER}, not {w.shape[-1]}")
    wg = _grouped(w, g.group)
    t = _flat_scale(w, GRIDS["iq4nl"], imp).double()
    d = _round_sig(_grouped(t.abs(), IQ4XS_SUPER // g.group).amax(-1) / IQ4XS_M_MAX, g.sig)
    unit = d.unsqueeze(-1).expand(*d.shape, IQ4XS_SUPER // g.group).flatten(-2)
    sign = torch.where(t < 0, -1.0, 1.0).float()
    m0 = (t.abs() / unit.double()).round()
    cands = [sign * (m0 + k).clamp(1, IQ4XS_M_MAX).float() for k in g.cands]
    return d, _search(wg, g, cands, unit, wt=_error_weight(wg, imp, g))


def _iq4xs_scale(d: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
    return d.unsqueeze(-1).expand(*d.shape, IQ4XS_SUPER // GRIDS["iq4xs"].group).flatten(-2) * m


def nbit_scale(w: torch.Tensor, spec: QuantSpec, importance: torch.Tensor | None = None) -> torch.Tensor:
    """The forward's signed per-group scale [..., in/g] in fp32 (int3: d·m per sub-block of 16; iq4xs: d·m per 32).
    D-57: `importance` [..., in] weights the search of the flat grids and IQ4_XS (Q3_K's search stays unweighted)."""
    g = GRIDS[spec.method]
    w = w.detach().float()
    if spec.method in Q3K_METHODS:
        return _q3k_scale(*_sliced(lambda x: _q3k_factors(x, g), w))
    if spec.method == "iq4xs":
        if importance is None:
            return _iq4xs_scale(*_sliced(lambda x: _iq4xs_factors(x, g), w))
        return _iq4xs_scale(*_sliced(lambda x, i: _iq4xs_factors(x, g, i), w, importance))
    if importance is None:
        return _sliced(lambda x: _flat_scale(x, g), w)
    return _sliced(lambda x, i: _flat_scale(x, g, i), w, importance)


def nbit_codes(w: torch.Tensor, s: torch.Tensor, spec: QuantSpec) -> torch.Tensor:
    """int8 codes on the format's grid, shape of w, for the signed scales s [..., in/g]."""
    g = GRIDS[spec.method]
    return _sliced(lambda x, y: _grid_codes(_grouped(x, g.group), y, g).flatten(-2).to(torch.int8),
                   w.detach().float(), s)


class _NBitSTE(torch.autograd.Function):
    """D-51: forward s·q on the format's grid, exactly; backward the plain STE (the identity to the latent)."""

    @staticmethod
    def forward(ctx, w, spec: QuantSpec, importance=None):
        s = nbit_scale(w, spec, importance)
        q = nbit_codes(w, s, spec)
        v = _code_values(q.to(torch.float32), GRIDS[spec.method])
        w_hat = (_grouped(v, spec.group_size) * s.unsqueeze(-1)).flatten(-2)
        return w_hat.to(w.dtype)

    @staticmethod
    def backward(ctx, g):
        return g, None, None


def pack_4bit(u: torch.Tensor) -> torch.Tensor:
    """uint8 values 0..15 [..., n] -> uint8 [..., n/2]; element 2i in the low nibble of byte i, 2i+1 in the high."""
    u = u.to(torch.uint8).unflatten(-1, (-1, 2))
    return u[..., 0] | (u[..., 1] << 4)


def unpack_4bit(packed: torch.Tensor) -> torch.Tensor:
    return torch.stack([packed & 15, packed >> 4], -1).flatten(-2)


def export_packed_nbit(w: torch.Tensor, spec: QuantSpec, importance: torch.Tensor | None = None) -> dict:
    """Packed N-bit tensor: codes (int2: 2-bit as `pack_2bit`, q+1; int3/int4: 4-bit, q - lo) and fp16 scales
    exactly as the forward used them; int3 adds `sub_scales` (int8 m), its `scales` being Q3_K's d."""
    g = GRIDS[spec.method]
    w = w.detach().float()
    out = {"method": spec.method, "group_size": g.group, "shape": tuple(w.shape)}
    if spec.method in Q3K_METHODS:
        d, m = _sliced(lambda x: _q3k_factors(x, g), w)
        s = _q3k_scale(d, m)
        out.update(scales=d.to(torch.float16), sub_scales=m.to(torch.int8))
    elif spec.method == "iq4xs":
        if importance is None:
            d, m = _sliced(lambda x: _iq4xs_factors(x, g), w)
        else:
            d, m = _sliced(lambda x, i: _iq4xs_factors(x, g, i), w, importance)
        s = _iq4xs_scale(d, m)
        out.update(scales=d.to(torch.float16), sub_scales=m.to(torch.int8))
    else:
        s = nbit_scale(w, spec, importance)
        out["scales"] = s.to(torch.float16)
    q = nbit_codes(w, s, spec)
    out["codes"] = (pack_2bit(q) if spec.method in INT2_METHODS
                    else pack_4bit((q.to(torch.int16) - g.lo).to(torch.uint8)))
    return out


def dequant_packed_nbit(p: dict, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
    g = GRIDS[p["method"]]
    if p["method"] in INT2_METHODS:
        q = unpack_2bit(p["codes"]).to(torch.float32)
    else:
        q = _code_values((unpack_4bit(p["codes"]).to(torch.int16) + g.lo).to(torch.float32), g)
    if p["method"] in Q3K_METHODS:
        s = _q3k_scale(p["scales"].to(torch.float32), p["sub_scales"].to(torch.float32))
    elif p["method"] == "iq4xs":
        s = _iq4xs_scale(p["scales"].to(torch.float32), p["sub_scales"].to(torch.float32))
    else:
        s = p["scales"].to(torch.float32)
    return (_grouped(q, g.group) * s.unsqueeze(-1)).flatten(-2).to(dtype)


# ---------------------------------------------------------------- diagnostics


def zero_fraction(codes: torch.Tensor) -> float:
    return (codes == 0).float().mean().item()


def flip_rate(old: torch.Tensor, new: torch.Tensor) -> float:
    return (old != new).float().mean().item()


# ---------------------------------------------------------------- slow references (tests only)


def reference_codes_and_scale(w: torch.Tensor, spec: QuantSpec) -> tuple[torch.Tensor, torch.Tensor]:
    """Element-by-element reference for absmean and twn over a 2-D tensor."""
    rows, cols = w.shape
    g = spec.group_size
    codes = torch.zeros(rows, cols, dtype=torch.int8)
    scales = torch.zeros(rows, cols // g, dtype=torch.float32)
    for r in range(rows):
        for j in range(cols // g):
            grp = [float(x) for x in w[r, j * g:(j + 1) * g]]
            mean_abs = sum(abs(x) for x in grp) / g
            if spec.method == "twn":
                delta = TWN_DELTA * mean_abs
                kept = [abs(x) for x in grp if abs(x) > delta]
                s = sum(kept) / max(len(kept), 1)
                q = [(1 if x > 0 else -1) if abs(x) > delta else 0 for x in grp]
            else:
                s = mean_abs
                q = None
            s = max(s, spec.eps)
            s = float(torch.tensor(s, dtype=torch.float32).to(torch.bfloat16).float())
            if q is None:
                # round-half-even of x/s clamped to [-1, 1], in exact arithmetic (s/2 is exact)
                q = [(1 if x > 0 else -1) if abs(x) > 0.5 * s else 0 for x in grp]
            scales[r, j] = s
            codes[r, j * g:(j + 1) * g] = torch.tensor(q, dtype=torch.int8)
    return codes, scales