Upload 3 files

pipeline.py

@@ -931,8 +931,9 @@ class Q2KQuantizer:
                 if seg_perm is not None:
                     from .segmentation import band_labels
 
-                    Xh = np.stack([b.dequantize() for b in res.blocks]).astype(np.float64)
-                    labels_slot = np.take_along_axis(seg_labels, seg_perm, axis=1)
+                    n_chk = min(n_blocks, 256)
+                    Xh = np.stack([b.dequantize() for b in res.blocks[:n_chk]]).astype(np.float64)
+                    labels_slot = np.take_along_axis(seg_labels[:n_chk], seg_perm[:n_chk], axis=1)
                     labels_rec = band_labels(Xh, max(2, int(cfg.segment_k)))
                     res.stats.segment_rounds = seg_rounds
                     res.stats.segment_flips = int(np.count_nonzero(labels_rec != labels_slot))
@@ -1262,8 +1263,9 @@ class Q2KQuantizer:
             # mis-derive from the stored slot values
             from .segmentation import band_labels
 
-            Xh = np.stack([b.dequantize() for b in blocks]).astype(np.float64)
-            labels_slot = np.take_along_axis(seg_labels, seg_perm, axis=1)
+            n_chk = min(n_blocks, 256)
+            Xh = np.stack([b.dequantize() for b in blocks[:n_chk]]).astype(np.float64)
+            labels_slot = np.take_along_axis(seg_labels[:n_chk], seg_perm[:n_chk], axis=1)
             labels_rec = band_labels(Xh, max(2, int(cfg.segment_k)))
             seg_flips = int(np.count_nonzero(labels_rec != labels_slot))
 

q3_k.py

@@ -1064,10 +1064,12 @@ def quantize_tensor_q3_k(
     rmse = float(np.sqrt(total_sse / x.size))
 
     if cfg.segment:
-        xh = np.asarray(dequantize_q3_k_bytes(b"".join(chunks), n_blocks), dtype=np.float64).reshape(
-            n_blocks, QK3_K
+        n_chk = min(n_blocks, 256)
+        raw_chk = b"".join(chunks)[: n_chk * Q3_K_BLOCK_BYTES]
+        xh = np.asarray(dequantize_q3_k_bytes(raw_chk, n_chk), dtype=np.float64).reshape(
+            n_chk, QK3_K
         )
-        labels_slot = np.take_along_axis(seg_labels, seg_order, axis=1)
+        labels_slot = np.take_along_axis(seg_labels[:n_chk], seg_order[:n_chk], axis=1)
         labels_rec = band_labels(xh, max(2, int(cfg.segment_k)))
         seg_flips = int(np.count_nonzero(labels_rec != labels_slot))
         _log.info(