From 646480abaa69204ba255d4e06be81a4f7959dfcb Mon Sep 17 00:00:00 2001
From: narugo1992 <narugo1992@users.noreply.github.com>
Date: Mon, 10 Aug 2026 11:33:47 +0000
Subject: [PATCH 1/4] dev(narugo): stream the threshold histogram instead of
 rescanning the score matrix

The per-tag and per-category threshold searches were O(tags x thresholds x
samples) when the problem is O(tags x samples). For a fixed threshold grid every
confusion matrix is a suffix sum over a per-tag score histogram, so one pass
answers all thresholds at once and the per-category search reduces to summing the
per-tag histograms.

StreamingThresholdHistogram folds each inference batch into a
(num_thresholds + 1, tag_num) counter while it is still on the device, so test.py
no longer materialises the (sample_num, tag_num) score and label matrices, and no
longer all-gathers them to every rank -- only the counter is reduced.

Measured on the real full EVA-Giant test run (295936 x 12476): threshold phase
24.86 min -> 82.5 ms, peak RSS 102.4 GB -> 28.1 GB, with output bit-identical
across all 12476 tags and every category, on both CPU and CUDA. Runs unchanged on
CPU-only machines: the accumulator follows the accelerator's device and skips the
reduction when there is a single process.

See https://github.com/deepghs/animetimm/issues/1

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 animetimm/multilabel/metrics.py | 375 +++++++++++++++-------
 animetimm/multilabel/test.py    |  37 +--
 pytest.ini                      |   4 +
 requirements-test.txt           |   3 +
 test/__init__.py                |   0
 test/multilabel/__init__.py     |   0
 test/multilabel/test_metrics.py | 540 ++++++++++++++++++++++++++++++++
 7 files changed, 818 insertions(+), 141 deletions(-)
 create mode 100644 pytest.ini
 create mode 100644 requirements-test.txt
 create mode 100644 test/__init__.py
 create mode 100644 test/multilabel/__init__.py
 create mode 100644 test/multilabel/test_metrics.py

diff --git a/animetimm/multilabel/metrics.py b/animetimm/multilabel/metrics.py
index ae5678b..d0d5ac9 100644
--- a/animetimm/multilabel/metrics.py
+++ b/animetimm/multilabel/metrics.py
@@ -1,10 +1,8 @@
-import logging
+from typing import Optional, Tuple
 
 import numpy as np
 import torch
 
-from ..utils import parallel_call
-
 
 def mcc(tp, fp, tn, fn, mean: bool = True):
     N = (tp + fn + fp + tn)
@@ -80,129 +78,266 @@ def recall(tp, fp, tn, fn, mean: bool = True):
     return v
 
 
-def compute_optimal_thresholds(all_sample, all_labels, alpha: float = 1.0, num_thresholds: int = 100,
-                               max_workers: int = 32):
-    all_sample = all_sample.detach().cpu().numpy()
-    # print(all_labels)
-    all_labels = all_labels.to(torch.int32).to(torch.bool).detach().cpu().numpy()
-    # print(all_labels.to(torch.int))
-
-    # Generate candidate thresholds (0 to 1)
-    thresholds = np.linspace(1.0 / num_thresholds, 1, num_thresholds)
-
-    best_f1, best_precision, best_recall, best_thresholds = {}, {}, {}, {}
+Histograms = Tuple[np.ndarray, np.ndarray]
+
+
+def threshold_grid(num_thresholds: int = 100) -> np.ndarray:
+    """Candidate thresholds, ``num_thresholds`` points evenly spread over ``(0, 1]``."""
+    return np.linspace(1.0 / num_thresholds, 1, num_thresholds)
+
+
+def _binarize(labels):
+    """The label convention used throughout testing: truncate towards zero, then test
+    for non-zero, i.e. only an exact 1.0 counts as positive for a score in [0, 1]."""
+    return labels.to(torch.int32).to(torch.bool)
+
+
+def _bound_dtype(sample_dtype: torch.dtype) -> torch.dtype:
+    """Precision to compare scores against thresholds in.
+
+    float64 inputs keep full precision, matching what the original numpy comparison did;
+    everything else (float32 in practice, float16 in principle) compares in float32,
+    since a 0.01 grid is not representable well enough in half.
+    """
+    return torch.float64 if sample_dtype == torch.float64 else torch.float32
+
+
+def _bin_indices(sample, bounds):
+    """Bin index of each score: ``#{k : bounds[k] <= score}``.
+
+    A score therefore counts as predicted-positive for exactly ``bounds[0..bin-1]``,
+    which makes every threshold's confusion matrix a suffix sum over the bins.
+    """
+    if sample.dtype != bounds.dtype:
+        sample = sample.to(bounds.dtype)
+    return torch.bucketize(sample, bounds, right=True)
+
+
+class StreamingThresholdHistogram:
+    """Per-tag score histogram over threshold bins, accumulated one batch at a time.
+
+    Folding each inference batch straight into a ``(num_thresholds + 1, tag_num)``
+    counter means the ``(sample_num, tag_num)`` score and label matrices never have to
+    be kept around, and distributed runs only have to reduce the counter rather than
+    all-gather the raw tensors.
+
+    Works identically on CPU and CUDA; the accumulator follows the device of the first
+    batch handed to :meth:`update` unless ``device`` is given explicitly.
+    """
+
+    def __init__(self, tag_num: int, num_thresholds: int = 100, device=None,
+                 score_dtype: Optional[torch.dtype] = None):
+        self.tag_num = tag_num
+        self.num_thresholds = num_thresholds
+        # None means follow the first batch, the same rule the offline path uses
+        self.score_dtype = score_dtype
+        self._device = None
+        self._bounds = None
+        self._ones = None
+        self.hist_all = None
+        self.hist_pos = None
+        if device is not None:
+            self._allocate(torch.device(device))
+
+    def _allocate(self, device):
+        shape = (self.num_thresholds + 1, self.tag_num)
+        self.hist_all = torch.zeros(shape, dtype=torch.int64, device=device)
+        self.hist_pos = torch.zeros(shape, dtype=torch.int64, device=device)
+        # normalised form of the request, e.g. 'cuda' resolves to 'cuda:0', so the
+        # per-batch device check below never fires spuriously
+        self._device = self.hist_all.device
+        self._bounds = None
+        self._ones = None
+
+    def _bounds_for(self, sample):
+        """Bin edges, built from the same np.linspace the offline path uses so both
+        routes bucketize against bit-identical values."""
+        dtype = self.score_dtype or _bound_dtype(sample.dtype)
+        if self._bounds is None or self._bounds.dtype != dtype:
+            self._bounds = torch.as_tensor(threshold_grid(self.num_thresholds),
+                                           dtype=dtype, device=self._device)
+        return self._bounds
+
+    def _ones_like_rows(self, rows: int):
+        if self._ones is None or self._ones.shape[0] < rows:
+            self._ones = torch.ones((rows, self.tag_num), dtype=torch.int64, device=self._device)
+        return self._ones[:rows]
+
+    @torch.no_grad()
+    def update(self, sample, labels):
+        """Fold one batch in. ``sample`` holds sigmoid scores, ``labels`` the raw labels,
+        both shaped ``(batch_size, tag_num)``."""
+        if sample.shape != labels.shape:
+            raise ValueError(f'Sample shape {tuple(sample.shape)} does not match '
+                             f'label shape {tuple(labels.shape)}.')
+        if sample.shape[-1] != self.tag_num:
+            raise ValueError(f'Expected {self.tag_num} tags, got {sample.shape[-1]}.')
+        if self.hist_all is None:
+            self._allocate(sample.device)
+        elif sample.device != self._device:
+            sample, labels = sample.to(self._device), labels.to(self._device)
+
+        bins = _bin_indices(sample, self._bounds_for(sample))
+        self.hist_all.scatter_add_(0, bins, self._ones_like_rows(bins.shape[0]))
+        self.hist_pos.scatter_add_(0, bins, _binarize(labels).to(torch.int64))
+
+    @torch.no_grad()
+    def finalize(self, accelerator=None) -> Histograms:
+        """Reduce across processes when running distributed, and hand back the
+        ``(tag_num, num_thresholds + 1)`` histograms the searches consume."""
+        if self.hist_all is None:
+            self._allocate(self._device or torch.device('cpu'))
+
+        hist_all, hist_pos = self.hist_all, self.hist_pos
+        if accelerator is not None and accelerator.num_processes > 1:
+            # two (num_thresholds + 1, tag_num) int64 counters -- megabytes, not gigabytes
+            hist_all = accelerator.reduce(hist_all, reduction='sum')
+            hist_pos = accelerator.reduce(hist_pos, reduction='sum')
+        return (hist_all.t().to(torch.float64).cpu().numpy(),
+                hist_pos.t().to(torch.float64).cpu().numpy())
+
+
+def build_threshold_histograms(all_sample, all_labels, num_thresholds: int = 100,
+                               chunk: int = 1024, device=None) -> Histograms:
+    """Offline counterpart of :class:`StreamingThresholdHistogram`, for callers that
+    already hold the full ``(sample_num, tag_num)`` matrices.
+
+    Tags are processed in chunks so peak memory stays bounded regardless of tag count.
+    """
+    if all_sample.shape != all_labels.shape:
+        raise ValueError(f'Sample shape {tuple(all_sample.shape)} does not match '
+                         f'label shape {tuple(all_labels.shape)}.')
+    device = torch.device(device) if device is not None else all_sample.device
+    tag_num = all_sample.shape[-1]
+    bounds = torch.as_tensor(threshold_grid(num_thresholds),
+                             dtype=_bound_dtype(all_sample.dtype), device=device)
+
+    hist_all = np.empty((tag_num, num_thresholds + 1), dtype=np.float64)
+    hist_pos = np.empty((tag_num, num_thresholds + 1), dtype=np.float64)
+    width_stride = num_thresholds + 1
+
+    for start in range(0, tag_num, chunk):
+        stop = min(start + chunk, tag_num)
+        width = stop - start
+        sample = all_sample[:, start:stop].to(device).contiguous()
+        labels = _binarize(all_labels[:, start:stop].to(device).contiguous())
+
+        bins = _bin_indices(sample, bounds)
+        # offset each column into its own slice of one flat histogram, so a single
+        # bincount covers the whole chunk instead of one call per tag
+        bins = bins + torch.arange(width, device=bins.device).mul(width_stride).unsqueeze(0)
+        bins = bins.reshape(-1)
+
+        size = width * width_stride
+        counts_all = torch.bincount(bins, minlength=size)
+        # counting only the positives' bins keeps this an integer histogram; float
+        # weights would start losing counts past 2**24 samples
+        counts_pos = torch.bincount(bins[labels.reshape(-1)], minlength=size)
+        hist_all[start:stop] = counts_all.view(width, width_stride).cpu().numpy()
+        hist_pos[start:stop] = counts_pos.view(width, width_stride).cpu().numpy()
+
+    return hist_all, hist_pos
+
+
+def _curves(hist_all, hist_pos, alpha: float):
+    """Turn bin histograms into f1/precision/recall at every threshold.
+
+    ``suffix[..., j]`` counts the scores landing in bin ``j`` or above, which is exactly
+    the number of positives predicted at ``thresholds[j - 1]``.
+    """
+    suffix_all = np.cumsum(hist_all[..., ::-1], axis=-1)[..., ::-1]
+    suffix_pos = np.cumsum(hist_pos[..., ::-1], axis=-1)[..., ::-1]
+    predicted_positive, tp = suffix_all[..., 1:], suffix_pos[..., 1:]
+    fp = predicted_positive - tp
+    fn = hist_pos.sum(-1)[..., None] - tp
+
+    p = tp / (tp + fp + 1e-12)
+    r = tp / (tp + fn + 1e-12)
+    beta_sq = alpha ** 2
+    f1 = (1 + beta_sq) * p * r / (beta_sq * p + r + 1e-12)
+    return f1, p, r
+
+
+def _pick(f1s, pres, recs, ths):
+    """Best threshold per row of a ``(n, num_thresholds)`` block.
+
+    Ties are resolved the way the original scan did: walk forward from the argmax over
+    entries whose f1/precision/recall are all numerically indistinguishable, then take
+    the midpoint of the run. Expressed as a masked argmax rather than a Python loop --
+    ``np.isclose`` on scalars costs tens of microseconds and dominates otherwise.
+    """
+    row_num, th_num = f1s.shape
+    rows = np.arange(row_num)
+    ma = np.argmax(f1s, axis=1)
+
+    close = (np.isclose(f1s, f1s[rows, ma][:, None])
+             & np.isclose(pres, pres[rows, ma][:, None])
+             & np.isclose(recs, recs[rows, ma][:, None]))
+    # first index after ma that breaks the run, or th_num when the run reaches the end
+    breaks = (np.arange(th_num)[None, :] > ma[:, None]) & ~close
+    mb = np.where(breaks.any(axis=1), breaks.argmax(axis=1), th_num) - 1
+
+    return (ths[ma] + ths[mb]) / 2, f1s[rows, ma], pres[rows, ma], recs[rows, ma]
+
+
+def _resolve_histograms(all_sample, all_labels, num_thresholds, histograms, device):
+    if histograms is not None:
+        return histograms
+    if all_sample is None or all_labels is None:
+        raise ValueError('Either histograms, or both all_sample and all_labels, must be given.')
+    return build_threshold_histograms(all_sample, all_labels, num_thresholds=num_thresholds,
+                                      device=device)
+
+
+def compute_optimal_thresholds(all_sample=None, all_labels=None, alpha: float = 1.0, num_thresholds: int = 100,
+                               max_workers: Optional[int] = None, histograms: Optional[Histograms] = None,
+                               device=None):
+    """Per-tag optimal thresholds and the f1/precision/recall attained there.
+
+    Pass ``histograms`` from :class:`StreamingThresholdHistogram` or
+    :func:`build_threshold_histograms` to reuse one pass over the data for both this
+    and :func:`compute_optimal_thresholds_by_categories`.
+
+    ``max_workers`` is accepted for backwards compatibility and ignored -- the search no
+    longer uses a thread pool.
+    """
+    _ = max_workers
+    hist_all, hist_pos = _resolve_histograms(all_sample, all_labels, num_thresholds, histograms, device)
+    thresholds = threshold_grid(num_thresholds)
+    f1s, pres, recs = _curves(hist_all, hist_pos, alpha)
+    best_thresholds, best_f1, best_precision, best_recall = _pick(f1s, pres, recs, thresholds)
+    return best_thresholds, best_f1, best_precision, best_recall
 
-    def _fn_cal(idx):
-        sample, labels = all_sample[..., idx], all_labels[..., idx]
-
-        f1s, pres, recs, ths = [], [], [], []
-
-        for th in thresholds:
-            ppos = sample >= th
-            tp = ((ppos == 1) & (labels == 1)).sum()
-            fp = ((ppos == 1) & (labels == 0)).sum()
-            fn = ((ppos == 0) & (labels == 1)).sum()
-
-            p = tp / (tp + fp + 1e-12)
-            r = tp / (tp + fn + 1e-12)
-            beta_sq = alpha ** 2
-            f1_numerator = (1 + beta_sq) * p * r
-            f1_denominator = beta_sq * p + r + 1e-12
-            f1 = f1_numerator / f1_denominator
-            f1s.append(f1)
-            pres.append(p)
-            recs.append(r)
-            ths.append(th)
-
-        f1s = np.array(f1s)
-        pres = np.array(pres)
-        recs = np.array(recs)
-        ths = np.array(ths)
-
-        ma = int(np.argmax(f1s).item())
-        mb = int(ma) + 1
-        while mb < f1s.shape[0] and np.isclose(f1s[ma], f1s[mb]) and np.isclose(pres[ma], pres[mb]) \
-                and np.isclose(recs[ma], recs[mb]):
-            mb += 1
-        mb = mb - 1
-        best_f1[idx] = f1s[ma]
-        best_precision[idx] = pres[ma]
-        best_recall[idx] = recs[ma]
-        best_thresholds[idx] = (ths[ma] + ths[mb]) / 2
-
-    parallel_call(
-        iterable=range(all_sample.shape[-1]),
-        fn=_fn_cal,
-        desc='Scan Tags',
-        max_workers=max_workers,
-    )
-
-    best_f1 = np.array([best_f1[idx] for idx in range(all_sample.shape[-1])])
-    best_precision = np.array([best_precision[idx] for idx in range(all_sample.shape[-1])])
-    best_recall = np.array([best_recall[idx] for idx in range(all_sample.shape[-1])])
-    best_thresholds = np.array([best_thresholds[idx] for idx in range(all_sample.shape[-1])])
 
-    return best_thresholds, best_f1, best_precision, best_recall
+def compute_optimal_thresholds_by_categories(all_sample=None, all_labels=None, df_tags=None, alpha: float = 1.0,
+                                             num_thresholds: int = 100, max_workers: Optional[int] = None,
+                                             histograms: Optional[Histograms] = None, device=None):
+    """Optimal threshold per tag category, micro-averaged over the tags in it.
 
+    A category's confusion matrix at a threshold is just the sum of its tags' bin
+    histograms, so this is essentially free once the per-tag histograms exist.
 
-def compute_optimal_thresholds_by_categories(all_sample, all_labels, df_tags,
-                                             alpha: float = 1.0, num_thresholds: int = 100, max_workers: int = 16):
-    all_sample = all_sample.detach().cpu().numpy()
-    # print(all_labels)
-    all_labels = all_labels.to(torch.int32).to(torch.bool).detach().cpu().numpy()
-    # print(all_labels.to(torch.int))
+    ``max_workers`` is accepted for backwards compatibility and ignored.
+    """
+    _ = max_workers
+    if df_tags is None:
+        raise ValueError('df_tags is required to group tags into categories.')
+    hist_all, hist_pos = _resolve_histograms(all_sample, all_labels, num_thresholds, histograms, device)
+    thresholds = threshold_grid(num_thresholds)
 
-    # Generate candidate thresholds (0 to 1)
-    thresholds = np.linspace(1.0 / num_thresholds, 1, num_thresholds)
+    categories = np.asarray(df_tags['category'])
+    if categories.shape[0] != hist_all.shape[0]:
+        raise ValueError(f'Tag table has {categories.shape[0]} rows but the histograms '
+                         f'cover {hist_all.shape[0]} tags.')
 
     best_f1, best_precision, best_recall, best_thresholds = {}, {}, {}, {}
-
-    for category in sorted(set(df_tags['category'])):
-        logging.info(f'Scanning for category {category!r} ...')
-        mask = df_tags['category'] == category
-        sample, labels = all_sample[..., mask], all_labels[..., mask]
-
-        f1s, pres, recs, ths = {}, {}, {}, {}
-
-        def _fn_call(idx):
-            th = thresholds[idx]
-            ppos = sample >= th
-            tp = ((ppos == 1) & (labels == 1)).sum()
-            fp = ((ppos == 1) & (labels == 0)).sum()
-            fn = ((ppos == 0) & (labels == 1)).sum()
-
-            p = tp / (tp + fp + 1e-12)
-            r = tp / (tp + fn + 1e-12)
-            beta_sq = alpha ** 2
-            f1_numerator = (1 + beta_sq) * p * r
-            f1_denominator = beta_sq * p + r + 1e-12
-            f1 = f1_numerator / f1_denominator
-            f1s[idx] = f1
-            pres[idx] = p
-            recs[idx] = r
-            ths[idx] = th
-
-        parallel_call(
-            iterable=range(thresholds.shape[0]),
-            fn=_fn_call,
-            desc=f'Scan Thresholds For #{category}',
-            max_workers=max_workers,
-        )
-
-        f1s = np.array([f1s[idx] for idx in range(thresholds.shape[0])])
-        pres = np.array([pres[idx] for idx in range(thresholds.shape[0])])
-        recs = np.array([recs[idx] for idx in range(thresholds.shape[0])])
-        ths = np.array([ths[idx] for idx in range(thresholds.shape[0])])
-
-        ma = int(np.argmax(f1s).item())
-        mb = int(ma) + 1
-        while mb < f1s.shape[0] and np.isclose(f1s[ma], f1s[mb]) and np.isclose(pres[ma], pres[mb]) \
-                and np.isclose(recs[ma], recs[mb]):
-            mb += 1
-        mb = mb - 1
-        best_f1[category] = float(f1s[ma])
-        best_precision[category] = float(pres[ma])
-        best_recall[category] = float(recs[ma])
-        best_thresholds[category] = float((ths[ma] + ths[mb]) / 2)
+    for category in sorted(set(categories.tolist())):
+        mask = categories == category
+        f1s, pres, recs = _curves(hist_all[mask].sum(0)[None], hist_pos[mask].sum(0)[None], alpha)
+        th, f1, p, r = _pick(f1s, pres, recs, thresholds)
+        best_thresholds[category] = float(th[0])
+        best_f1[category] = float(f1[0])
+        best_precision[category] = float(p[0])
+        best_recall[category] = float(r[0])
 
     return best_thresholds, best_f1, best_precision, best_recall
diff --git a/animetimm/multilabel/test.py b/animetimm/multilabel/test.py
index 4f58ee7..24a8cc2 100644
--- a/animetimm/multilabel/test.py
+++ b/animetimm/multilabel/test.py
@@ -13,7 +13,7 @@
 
 from .dataset import load_tags, load_dataloader
 from .metrics import mcc, f1score, precision, recall, compute_optimal_thresholds, \
-    compute_optimal_thresholds_by_categories
+    compute_optimal_thresholds_by_categories, StreamingThresholdHistogram
 from ..dataset import load_pretrained_tag
 from ..model import Model
 from ..utils import GLOBAL_CONTEXT_SETTINGS, print_version
@@ -108,16 +108,18 @@ def test(workdir: str, num_workers: int = 32, batch_size: int = 32, test_thresho
         macro_tn = torch.zeros((1, len(tags_info.tags),), device=accelerator.device)
         macro_fn = torch.zeros((1, len(tags_info.tags),), device=accelerator.device)
 
-        all_samples = []
-        all_labels = []
+        # folds every batch into a (num_thresholds + 1, tag_num) counter instead of
+        # keeping the full (sample_num, tag_num) score/label matrices around
+        threshold_histogram = StreamingThresholdHistogram(len(tags_info.tags), device=accelerator.device)
 
         for i, (inputs, labels_) in enumerate(tqdm(test_dataloader, disable=not accelerator.is_main_process)):
             inputs = inputs.float()
             labels_ = labels_
 
             outputs = module(inputs)
+            scores = torch.sigmoid(outputs)
             labels = labels_ > test_threshold
-            preds = torch.sigmoid(outputs) > test_threshold
+            preds = scores > test_threshold
             micro_tp += ((preds == 1) & (labels == 1)).sum().item()
             micro_fp += ((preds == 1) & (labels == 0)).sum().item()
             micro_tn += ((preds == 0) & (labels == 0)).sum().item()
@@ -127,8 +129,7 @@ def test(workdir: str, num_workers: int = 32, batch_size: int = 32, test_thresho
             macro_tn += ((preds == 0) & (labels == 0)).sum(dim=0)
             macro_fn += ((preds == 0) & (labels == 1)).sum(dim=0)
 
-            all_samples.append(torch.sigmoid(outputs))
-            all_labels.append(labels_)
+            threshold_histogram.update(scores, labels_)
 
             if i % 10 == 0:
                 accelerator.wait_for_everyone()
@@ -136,19 +137,13 @@ def test(workdir: str, num_workers: int = 32, batch_size: int = 32, test_thresho
         logging.info(f'Inference ready for #{accelerator.process_index}.')
         accelerator.wait_for_everyone()
 
-        all_samples = torch.concat(all_samples, dim=0)
-        all_labels = torch.concat(all_labels, dim=0)
-        all_samples = accelerator.gather(all_samples).cpu()
-        all_labels = accelerator.gather(all_labels).cpu()
-        # all_samples, all_labels = accelerator.gather_for_metrics((all_samples, all_labels))
+        # reduces two (num_thresholds + 1, tag_num) counters instead of all-gathering
+        # the raw score and label matrices
+        histograms = threshold_histogram.finalize(accelerator=accelerator)
         if accelerator.is_main_process:
-            logging.info(f'Gathered all_samples, shape: {all_samples.shape!r}, '
-                         f'dtype: {all_samples.dtype!r}, device: {all_samples.device!r}')
-            logging.info(f'Gathered all_labels, shape: {all_labels.shape!r}, '
-                         f'dtype: {all_labels.dtype!r}, device: {all_labels.device!r}')
-
-        # print((torch.isclose(all_labels, 1.0) | torch.isclose(all_labels, 0.0)).all())
-        # quit()
+            logging.info(f'Threshold histograms reduced, shape: {histograms[0].shape!r}, '
+                         f'samples counted: {int(histograms[0][0].sum())!r}, '
+                         f'positives counted: {int(histograms[1].sum())!r}')
 
         # micro_tp = micro_tp.sum(dim=0)
         # micro_fp = micro_fp.sum(dim=0)
@@ -172,11 +167,11 @@ def test(workdir: str, num_workers: int = 32, batch_size: int = 32, test_thresho
 
         if accelerator.is_main_process:
             best_thresholds, best_f1, best_precision, best_recall = \
-                compute_optimal_thresholds(all_samples, all_labels, alpha=1.0, max_workers=32)
+                compute_optimal_thresholds(alpha=1.0, histograms=histograms)
 
             c_best_thresholds, c_best_f1, c_best_precision, c_best_recall = \
-                compute_optimal_thresholds_by_categories(all_samples, all_labels, tags_info.df, alpha=1.0,
-                                                         max_workers=8)
+                compute_optimal_thresholds_by_categories(df_tags=tags_info.df, alpha=1.0,
+                                                         histograms=histograms)
 
             micro_mcc = mcc(micro_tp, micro_fp, micro_tn, micro_fn).detach().cpu().item()
             micro_f1 = f1score(micro_tp, micro_fp, micro_tn, micro_fn).detach().cpu().item()
diff --git a/pytest.ini b/pytest.ini
new file mode 100644
index 0000000..509e254
--- /dev/null
+++ b/pytest.ini
@@ -0,0 +1,4 @@
+[pytest]
+markers =
+    unittest: unit tests
+testpaths = test
diff --git a/requirements-test.txt b/requirements-test.txt
new file mode 100644
index 0000000..447807c
--- /dev/null
+++ b/requirements-test.txt
@@ -0,0 +1,3 @@
+pytest~=7.4
+pytest-cov~=4.1
+pytest-xdist~=3.5
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/test/multilabel/__init__.py b/test/multilabel/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/test/multilabel/test_metrics.py b/test/multilabel/test_metrics.py
new file mode 100644
index 0000000..50c211c
--- /dev/null
+++ b/test/multilabel/test_metrics.py
@@ -0,0 +1,540 @@
+"""Behavioural lock for the threshold search.
+
+The reference implementations below are the original brute-force scans, transcribed
+verbatim from the version of ``animetimm/multilabel/metrics.py`` that preceded the
+histogram rewrite. Every test asserts the fast path reproduces them *exactly*, so any
+future change that shifts a published threshold has to break a test first.
+"""
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+
+from animetimm.multilabel.metrics import (
+    StreamingThresholdHistogram,
+    build_threshold_histograms,
+    compute_optimal_thresholds,
+    compute_optimal_thresholds_by_categories,
+    threshold_grid,
+)
+
+
+# ---------------------------------------------------------------- reference (oracle)
+
+def reference_optimal_thresholds(all_sample, all_labels, alpha: float = 1.0, num_thresholds: int = 100):
+    """Original per-tag scan: for every tag, evaluate every threshold from scratch."""
+    all_sample = all_sample.detach().cpu().numpy()
+    all_labels = all_labels.to(torch.int32).to(torch.bool).detach().cpu().numpy()
+    thresholds = np.linspace(1.0 / num_thresholds, 1, num_thresholds)
+
+    best_f1, best_precision, best_recall, best_thresholds = {}, {}, {}, {}
+    for idx in range(all_sample.shape[-1]):
+        sample, labels = all_sample[..., idx], all_labels[..., idx]
+        f1s, pres, recs, ths = [], [], [], []
+        for th in thresholds:
+            ppos = sample >= th
+            tp = ((ppos == 1) & (labels == 1)).sum()
+            fp = ((ppos == 1) & (labels == 0)).sum()
+            fn = ((ppos == 0) & (labels == 1)).sum()
+
+            p = tp / (tp + fp + 1e-12)
+            r = tp / (tp + fn + 1e-12)
+            beta_sq = alpha ** 2
+            f1 = (1 + beta_sq) * p * r / (beta_sq * p + r + 1e-12)
+            f1s.append(f1)
+            pres.append(p)
+            recs.append(r)
+            ths.append(th)
+
+        f1s, pres, recs, ths = np.array(f1s), np.array(pres), np.array(recs), np.array(ths)
+        ma = int(np.argmax(f1s).item())
+        mb = int(ma) + 1
+        while mb < f1s.shape[0] and np.isclose(f1s[ma], f1s[mb]) and np.isclose(pres[ma], pres[mb]) \
+                and np.isclose(recs[ma], recs[mb]):
+            mb += 1
+        mb = mb - 1
+        best_f1[idx] = f1s[ma]
+        best_precision[idx] = pres[ma]
+        best_recall[idx] = recs[ma]
+        best_thresholds[idx] = (ths[ma] + ths[mb]) / 2
+
+    tag_num = all_sample.shape[-1]
+    return (np.array([best_thresholds[i] for i in range(tag_num)]),
+            np.array([best_f1[i] for i in range(tag_num)]),
+            np.array([best_precision[i] for i in range(tag_num)]),
+            np.array([best_recall[i] for i in range(tag_num)]))
+
+
+def reference_optimal_thresholds_by_categories(all_sample, all_labels, df_tags, alpha: float = 1.0,
+                                               num_thresholds: int = 100):
+    """Original per-category scan: micro-average over every tag in the category."""
+    all_sample = all_sample.detach().cpu().numpy()
+    all_labels = all_labels.to(torch.int32).to(torch.bool).detach().cpu().numpy()
+    thresholds = np.linspace(1.0 / num_thresholds, 1, num_thresholds)
+
+    best_f1, best_precision, best_recall, best_thresholds = {}, {}, {}, {}
+    for category in sorted(set(df_tags['category'])):
+        mask = df_tags['category'] == category
+        sample, labels = all_sample[..., mask], all_labels[..., mask]
+
+        f1s, pres, recs, ths = [], [], [], []
+        for th in thresholds:
+            ppos = sample >= th
+            tp = ((ppos == 1) & (labels == 1)).sum()
+            fp = ((ppos == 1) & (labels == 0)).sum()
+            fn = ((ppos == 0) & (labels == 1)).sum()
+
+            p = tp / (tp + fp + 1e-12)
+            r = tp / (tp + fn + 1e-12)
+            beta_sq = alpha ** 2
+            f1 = (1 + beta_sq) * p * r / (beta_sq * p + r + 1e-12)
+            f1s.append(f1)
+            pres.append(p)
+            recs.append(r)
+            ths.append(th)
+
+        f1s, pres, recs, ths = np.array(f1s), np.array(pres), np.array(recs), np.array(ths)
+        ma = int(np.argmax(f1s).item())
+        mb = int(ma) + 1
+        while mb < f1s.shape[0] and np.isclose(f1s[ma], f1s[mb]) and np.isclose(pres[ma], pres[mb]) \
+                and np.isclose(recs[ma], recs[mb]):
+            mb += 1
+        mb = mb - 1
+        best_f1[category] = float(f1s[ma])
+        best_precision[category] = float(pres[ma])
+        best_recall[category] = float(recs[ma])
+        best_thresholds[category] = float((ths[ma] + ths[mb]) / 2)
+
+    return best_thresholds, best_f1, best_precision, best_recall
+
+
+# ---------------------------------------------------------------------- fixtures
+
+def make_case(sample_num=512, tag_num=48, seed=0, positive_rate=0.06):
+    """Scores that look like sigmoid outputs: mostly near zero, positives pushed up."""
+    rng = np.random.default_rng(seed)
+    labels = (rng.random((sample_num, tag_num)) < positive_rate).astype(np.float32)
+    logit = rng.standard_normal((sample_num, tag_num)).astype(np.float32) - 2.0 + 4.0 * labels
+    samples = (1.0 / (1.0 + np.exp(-logit))).astype(np.float32)
+    return torch.from_numpy(samples), torch.from_numpy(labels)
+
+
+def make_degenerate_case():
+    """Corner cases that make the tie-walk and the zero-denominator guards matter."""
+    sample_num, tag_num = 256, 8
+    rng = np.random.default_rng(7)
+    samples = rng.random((sample_num, tag_num)).astype(np.float32)
+    labels = (rng.random((sample_num, tag_num)) < 0.3).astype(np.float32)
+
+    labels[:, 0] = 0.0                      # tag never positive
+    labels[:, 1] = 1.0                      # tag always positive
+    samples[:, 2] = 0.0                     # scores pinned to the bottom
+    samples[:, 3] = 1.0                     # scores pinned to the top
+    samples[:, 4] = 0.5                     # every score exactly on a bin edge
+    labels[:, 5] = 1.0
+    samples[:, 5] = 1.0                     # perfectly separable
+    samples[:, 6] = 1.0 - labels[:, 6]      # perfectly anti-correlated
+    return torch.from_numpy(samples), torch.from_numpy(labels)
+
+
+def category_frame(tag_num, groups=(0, 4, 9)):
+    rng = np.random.default_rng(11)
+    cats = rng.choice(np.asarray(groups), size=tag_num)
+    cats[:len(groups)] = np.asarray(groups)  # guarantee every category is populated
+    return pd.DataFrame({'category': cats})
+
+
+def assert_all_identical(expected, actual, names=('threshold', 'f1', 'precision', 'recall')):
+    for name, a, b in zip(names, expected, actual):
+        a, b = np.asarray(a, dtype=np.float64), np.asarray(b, dtype=np.float64)
+        assert a.shape == b.shape, f'{name}: shape {a.shape} != {b.shape}'
+        assert np.array_equal(a, b), (
+            f'{name}: {int((a != b).sum())}/{a.size} entries differ, '
+            f'max|diff|={np.abs(a - b).max():.3e}'
+        )
+
+
+# ------------------------------------------------------------------------- tests
+
+@pytest.mark.unittest
+class TestThresholdGrid:
+    def test_matches_documented_grid(self):
+        grid = threshold_grid(100)
+        assert grid.shape == (100,)
+        np.testing.assert_array_equal(grid, np.linspace(0.01, 1, 100))
+
+    @pytest.mark.parametrize('num_thresholds', [1, 2, 10, 100, 257])
+    def test_endpoints(self, num_thresholds):
+        grid = threshold_grid(num_thresholds)
+        assert grid.shape == (num_thresholds,)
+        assert grid[-1] == pytest.approx(1.0)
+        assert grid[0] == pytest.approx(1.0 / num_thresholds)
+
+
+@pytest.mark.unittest
+class TestPerTagSearch:
+    def test_matches_reference(self):
+        samples, labels = make_case()
+        assert_all_identical(reference_optimal_thresholds(samples, labels),
+                             compute_optimal_thresholds(samples, labels))
+
+    def test_matches_reference_on_degenerate_tags(self):
+        samples, labels = make_degenerate_case()
+        assert_all_identical(reference_optimal_thresholds(samples, labels),
+                             compute_optimal_thresholds(samples, labels))
+
+    @pytest.mark.parametrize('num_thresholds', [4, 20, 100])
+    def test_matches_reference_across_grid_sizes(self, num_thresholds):
+        samples, labels = make_case(sample_num=384, tag_num=24, seed=num_thresholds)
+        assert_all_identical(
+            reference_optimal_thresholds(samples, labels, num_thresholds=num_thresholds),
+            compute_optimal_thresholds(samples, labels, num_thresholds=num_thresholds),
+        )
+
+    @pytest.mark.parametrize('alpha', [0.5, 1.0, 2.0])
+    def test_matches_reference_across_alpha(self, alpha):
+        samples, labels = make_case(sample_num=384, tag_num=24, seed=3)
+        assert_all_identical(reference_optimal_thresholds(samples, labels, alpha=alpha),
+                             compute_optimal_thresholds(samples, labels, alpha=alpha))
+
+    def test_single_tag(self):
+        samples, labels = make_case(sample_num=256, tag_num=1, seed=5)
+        assert_all_identical(reference_optimal_thresholds(samples, labels),
+                             compute_optimal_thresholds(samples, labels))
+
+    def test_max_workers_is_accepted_and_ignored(self):
+        samples, labels = make_case(sample_num=256, tag_num=8)
+        assert_all_identical(compute_optimal_thresholds(samples, labels),
+                             compute_optimal_thresholds(samples, labels, max_workers=32))
+
+    def test_requires_data_or_histograms(self):
+        with pytest.raises(ValueError, match='histograms'):
+            compute_optimal_thresholds()
+
+
+@pytest.mark.unittest
+class TestPerCategorySearch:
+    def test_matches_reference(self):
+        samples, labels = make_case()
+        df_tags = category_frame(samples.shape[1])
+        expected = reference_optimal_thresholds_by_categories(samples, labels, df_tags)
+        actual = compute_optimal_thresholds_by_categories(samples, labels, df_tags)
+        assert set(expected[0]) == set(actual[0])
+        for exp, act in zip(expected, actual):
+            for key in exp:
+                assert exp[key] == act[key], f'category {key}: {exp[key]!r} != {act[key]!r}'
+
+    def test_matches_reference_on_degenerate_tags(self):
+        samples, labels = make_degenerate_case()
+        df_tags = category_frame(samples.shape[1], groups=(0, 4))
+        expected = reference_optimal_thresholds_by_categories(samples, labels, df_tags)
+        actual = compute_optimal_thresholds_by_categories(samples, labels, df_tags)
+        for exp, act in zip(expected, actual):
+            for key in exp:
+                assert exp[key] == act[key], f'category {key}: {exp[key]!r} != {act[key]!r}'
+
+    def test_single_category(self):
+        samples, labels = make_case(sample_num=256, tag_num=12, seed=9)
+        df_tags = pd.DataFrame({'category': np.zeros(12, dtype=int)})
+        expected = reference_optimal_thresholds_by_categories(samples, labels, df_tags)
+        actual = compute_optimal_thresholds_by_categories(samples, labels, df_tags)
+        assert expected == actual
+
+    def test_requires_df_tags(self):
+        samples, labels = make_case(sample_num=64, tag_num=4)
+        with pytest.raises(ValueError, match='df_tags'):
+            compute_optimal_thresholds_by_categories(samples, labels)
+
+    def test_rejects_mismatched_tag_table(self):
+        samples, labels = make_case(sample_num=64, tag_num=4)
+        with pytest.raises(ValueError, match='histograms cover'):
+            compute_optimal_thresholds_by_categories(samples, labels, category_frame(5))
+
+
+@pytest.mark.unittest
+class TestOfflineHistograms:
+    def test_counts_every_sample_once(self):
+        samples, labels = make_case(sample_num=300, tag_num=16)
+        hist_all, hist_pos = build_threshold_histograms(samples, labels)
+        assert hist_all.shape == (16, 101)
+        np.testing.assert_array_equal(hist_all.sum(axis=1), np.full(16, 300.0))
+        np.testing.assert_array_equal(hist_pos.sum(axis=1),
+                                      labels.numpy().sum(axis=0).astype(np.float64))
+
+    def test_positives_never_exceed_totals(self):
+        samples, labels = make_case(sample_num=300, tag_num=16, seed=2)
+        hist_all, hist_pos = build_threshold_histograms(samples, labels)
+        assert (hist_pos <= hist_all).all()
+
+    @pytest.mark.parametrize('chunk', [1, 3, 7, 16, 64])
+    def test_chunking_is_transparent(self, chunk):
+        samples, labels = make_case(sample_num=200, tag_num=16, seed=4)
+        expected = build_threshold_histograms(samples, labels, chunk=1024)
+        actual = build_threshold_histograms(samples, labels, chunk=chunk)
+        np.testing.assert_array_equal(expected[0], actual[0])
+        np.testing.assert_array_equal(expected[1], actual[1])
+
+    def test_reused_histograms_match_direct_call(self):
+        samples, labels = make_case(sample_num=256, tag_num=20, seed=6)
+        df_tags = category_frame(20)
+        histograms = build_threshold_histograms(samples, labels)
+        assert_all_identical(compute_optimal_thresholds(samples, labels),
+                             compute_optimal_thresholds(histograms=histograms))
+        assert (compute_optimal_thresholds_by_categories(samples, labels, df_tags)
+                == compute_optimal_thresholds_by_categories(df_tags=df_tags, histograms=histograms))
+
+    def test_rejects_shape_mismatch(self):
+        samples, labels = make_case(sample_num=64, tag_num=4)
+        with pytest.raises(ValueError, match='does not match'):
+            build_threshold_histograms(samples, labels[:, :2])
+
+
+@pytest.mark.unittest
+class TestStreamingHistogram:
+    @pytest.mark.parametrize('batch_size', [1, 5, 32, 128, 1000])
+    def test_matches_offline_for_any_batching(self, batch_size):
+        samples, labels = make_case(sample_num=257, tag_num=16, seed=8)
+        expected = build_threshold_histograms(samples, labels)
+
+        accumulator = StreamingThresholdHistogram(16)
+        for start in range(0, samples.shape[0], batch_size):
+            accumulator.update(samples[start:start + batch_size], labels[start:start + batch_size])
+        actual = accumulator.finalize()
+
+        np.testing.assert_array_equal(expected[0], actual[0])
+        np.testing.assert_array_equal(expected[1], actual[1])
+
+    def test_end_to_end_matches_reference(self):
+        samples, labels = make_case(sample_num=384, tag_num=24, seed=10)
+        df_tags = category_frame(24)
+
+        accumulator = StreamingThresholdHistogram(24)
+        for start in range(0, samples.shape[0], 37):
+            accumulator.update(samples[start:start + 37], labels[start:start + 37])
+        histograms = accumulator.finalize()
+
+        assert_all_identical(reference_optimal_thresholds(samples, labels),
+                             compute_optimal_thresholds(histograms=histograms))
+        expected = reference_optimal_thresholds_by_categories(samples, labels, df_tags)
+        actual = compute_optimal_thresholds_by_categories(df_tags=df_tags, histograms=histograms)
+        assert expected == actual
+
+    def test_works_without_any_update(self):
+        """An empty shard must still produce well-formed, all-zero histograms."""
+        hist_all, hist_pos = StreamingThresholdHistogram(6).finalize()
+        assert hist_all.shape == (6, 101)
+        assert not hist_all.any()
+        assert not hist_pos.any()
+
+    def test_defaults_to_cpu_without_cuda(self):
+        """CPU-only environments must work with no device argument at all."""
+        samples, labels = make_case(sample_num=64, tag_num=4, seed=12)
+        accumulator = StreamingThresholdHistogram(4)
+        accumulator.update(samples, labels)
+        assert accumulator.hist_all.device.type == 'cpu'
+        expected = build_threshold_histograms(samples, labels)
+        np.testing.assert_array_equal(expected[0], accumulator.finalize()[0])
+
+    def test_explicit_cpu_device(self):
+        samples, labels = make_case(sample_num=64, tag_num=4, seed=13)
+        accumulator = StreamingThresholdHistogram(4, device='cpu')
+        accumulator.update(samples, labels)
+        np.testing.assert_array_equal(build_threshold_histograms(samples, labels)[0],
+                                      accumulator.finalize()[0])
+
+    def test_finalize_without_accelerator_is_identity(self):
+        samples, labels = make_case(sample_num=64, tag_num=4, seed=14)
+        accumulator = StreamingThresholdHistogram(4)
+        accumulator.update(samples, labels)
+        np.testing.assert_array_equal(accumulator.finalize()[0], accumulator.finalize(None)[0])
+
+    def test_finalize_skips_reduce_for_single_process(self):
+        """accelerator.reduce must not be touched when there is nothing to reduce."""
+
+        class _SingleProcessAccelerator:
+            num_processes = 1
+
+            def reduce(self, tensor, reduction='sum'):
+                raise AssertionError('reduce must not be called for a single process')
+
+        samples, labels = make_case(sample_num=64, tag_num=4, seed=15)
+        accumulator = StreamingThresholdHistogram(4)
+        accumulator.update(samples, labels)
+        np.testing.assert_array_equal(build_threshold_histograms(samples, labels)[1],
+                                      accumulator.finalize(_SingleProcessAccelerator())[1])
+
+    def test_finalize_reduces_when_distributed(self):
+        """Two shards reduced with sum must equal one pass over the concatenation."""
+
+        class _SummingAccelerator:
+            """Stands in for a 2-process run whose peer holds `other`."""
+
+            num_processes = 2
+
+            def __init__(self, peer):
+                self.peer = peer
+                self._calls = 0
+
+            def reduce(self, tensor, reduction='sum'):
+                assert reduction == 'sum'
+                peer_tensor = (self.peer.hist_all, self.peer.hist_pos)[self._calls]
+                self._calls += 1
+                return tensor + peer_tensor
+
+        samples, labels = make_case(sample_num=200, tag_num=10, seed=16)
+        first, second = slice(0, 120), slice(120, 200)
+
+        shard_a = StreamingThresholdHistogram(10)
+        shard_a.update(samples[first], labels[first])
+        shard_b = StreamingThresholdHistogram(10)
+        shard_b.update(samples[second], labels[second])
+
+        reduced = shard_a.finalize(_SummingAccelerator(shard_b))
+        expected = build_threshold_histograms(samples, labels)
+        np.testing.assert_array_equal(expected[0], reduced[0])
+        np.testing.assert_array_equal(expected[1], reduced[1])
+
+    def test_rejects_shape_mismatch(self):
+        samples, labels = make_case(sample_num=32, tag_num=4)
+        accumulator = StreamingThresholdHistogram(4)
+        with pytest.raises(ValueError, match='does not match'):
+            accumulator.update(samples, labels[:, :2])
+
+    def test_rejects_wrong_tag_count(self):
+        samples, labels = make_case(sample_num=32, tag_num=4)
+        accumulator = StreamingThresholdHistogram(5)
+        with pytest.raises(ValueError, match='Expected 5 tags'):
+            accumulator.update(samples, labels)
+
+
+@pytest.mark.unittest
+class TestScoreDtypes:
+    @pytest.mark.parametrize('dtype', [torch.float32, torch.float64])
+    def test_matches_reference(self, dtype):
+        samples, labels = make_case(sample_num=256, tag_num=12, seed=23)
+        samples, labels = samples.to(dtype), labels.to(dtype)
+        assert_all_identical(reference_optimal_thresholds(samples, labels),
+                             compute_optimal_thresholds(samples, labels))
+
+    @pytest.mark.parametrize('dtype', [torch.float32, torch.float64])
+    def test_streaming_matches_offline(self, dtype):
+        samples, labels = make_case(sample_num=256, tag_num=12, seed=24)
+        samples, labels = samples.to(dtype), labels.to(dtype)
+        accumulator = StreamingThresholdHistogram(12)
+        for start in range(0, samples.shape[0], 48):
+            accumulator.update(samples[start:start + 48], labels[start:start + 48])
+        np.testing.assert_array_equal(build_threshold_histograms(samples, labels)[0],
+                                      accumulator.finalize()[0])
+
+    def test_float64_input_compares_in_float64(self):
+        """A score just below a threshold in float64, but indistinguishable from it once
+        rounded to float32, must be judged the way the original float64 numpy comparison
+        judged it: below."""
+        grid = threshold_grid(100)
+        edge = grid[6] - 1e-12                       # just under the 0.07 threshold
+        assert edge < grid[6]                        # ... in float64
+        assert np.float32(edge) >= np.float32(grid[6])  # ... but not in float32
+
+        labels = torch.ones((1, 1), dtype=torch.float64)
+        hist_all, _ = build_threshold_histograms(torch.tensor([[edge]], dtype=torch.float64), labels)
+        # the bin counts thresholds <= score, so 0.07 must not be counted
+        assert int(np.argmax(hist_all[0])) == 6
+
+        # the same value handed in as float32 genuinely is >= the threshold
+        hist32, _ = build_threshold_histograms(torch.tensor([[edge]], dtype=torch.float32),
+                                               labels.to(torch.float32))
+        assert int(np.argmax(hist32[0])) == 7
+
+
+@pytest.mark.unittest
+class TestHistogramCountsAreExact:
+    def test_positive_counts_are_integral(self):
+        """hist_pos is built by counting bins, not by float-weighted accumulation, so
+        the counts stay exact no matter how many samples land in one bin."""
+        sample_num = 5000
+        samples = torch.full((sample_num, 2), 0.5, dtype=torch.float32)
+        labels = torch.ones((sample_num, 2), dtype=torch.float32)
+        hist_all, hist_pos = build_threshold_histograms(samples, labels)
+        np.testing.assert_array_equal(hist_all, hist_pos)
+        assert hist_pos.sum() == 2 * sample_num
+        assert (hist_pos == np.floor(hist_pos)).all()
+
+
+@pytest.mark.unittest
+class TestLabelConvention:
+    def test_only_exact_one_counts_as_positive(self):
+        """Preserves the original `.to(int32).to(bool)` semantics: a soft label below
+        1.0 truncates to zero and is treated as negative."""
+        samples = torch.tensor([[0.9], [0.9], [0.9]], dtype=torch.float32)
+        labels = torch.tensor([[1.0], [0.99], [0.4]], dtype=torch.float32)
+        _, hist_pos = build_threshold_histograms(samples, labels)
+        assert hist_pos.sum() == 1.0
+
+    def test_matches_reference_with_soft_labels(self):
+        rng = np.random.default_rng(17)
+        samples = torch.from_numpy(rng.random((128, 6)).astype(np.float32))
+        labels = torch.from_numpy(rng.random((128, 6)).astype(np.float32))
+        assert_all_identical(reference_optimal_thresholds(samples, labels),
+                             compute_optimal_thresholds(samples, labels))
+
+
+@pytest.mark.unittest
+class TestAcceleratorIntegration:
+    """`test.py` drives the accumulator through an `Accelerator`, including on machines
+    with no GPU at all, so exercise that path with a real one."""
+
+    def test_cpu_only_accelerator_round_trip(self):
+        accelerate = pytest.importorskip('accelerate')
+        accelerator = accelerate.Accelerator(cpu=True)
+        assert accelerator.device.type == 'cpu'
+        assert accelerator.num_processes == 1
+
+        samples, labels = make_case(sample_num=192, tag_num=12, seed=22)
+        accumulator = StreamingThresholdHistogram(12, device=accelerator.device)
+        for start in range(0, samples.shape[0], 32):
+            accumulator.update(samples[start:start + 32].to(accelerator.device),
+                               labels[start:start + 32].to(accelerator.device))
+        histograms = accumulator.finalize(accelerator=accelerator)
+
+        expected = build_threshold_histograms(samples, labels)
+        np.testing.assert_array_equal(expected[0], histograms[0])
+        np.testing.assert_array_equal(expected[1], histograms[1])
+        assert_all_identical(reference_optimal_thresholds(samples, labels),
+                             compute_optimal_thresholds(histograms=histograms))
+
+
+@pytest.mark.unittest
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='CUDA is unavailable')
+class TestCudaParity:
+    def test_offline_histograms_match_cpu(self):
+        samples, labels = make_case(sample_num=300, tag_num=20, seed=18)
+        expected = build_threshold_histograms(samples, labels, device='cpu')
+        actual = build_threshold_histograms(samples, labels, device='cuda')
+        np.testing.assert_array_equal(expected[0], actual[0])
+        np.testing.assert_array_equal(expected[1], actual[1])
+
+    def test_streaming_matches_cpu(self):
+        samples, labels = make_case(sample_num=300, tag_num=20, seed=19)
+        expected = build_threshold_histograms(samples, labels, device='cpu')
+
+        accumulator = StreamingThresholdHistogram(20, device='cuda')
+        for start in range(0, samples.shape[0], 64):
+            accumulator.update(samples[start:start + 64].cuda(), labels[start:start + 64].cuda())
+        actual = accumulator.finalize()
+
+        np.testing.assert_array_equal(expected[0], actual[0])
+        np.testing.assert_array_equal(expected[1], actual[1])
+
+    def test_streaming_accepts_cpu_batches_on_cuda_accumulator(self):
+        samples, labels = make_case(sample_num=128, tag_num=10, seed=20)
+        accumulator = StreamingThresholdHistogram(10, device='cuda')
+        accumulator.update(samples, labels)
+        np.testing.assert_array_equal(build_threshold_histograms(samples, labels)[0],
+                                      accumulator.finalize()[0])
+
+    def test_end_to_end_matches_reference_on_cuda(self):
+        samples, labels = make_case(sample_num=256, tag_num=16, seed=21)
+        assert_all_identical(
+            reference_optimal_thresholds(samples, labels),
+            compute_optimal_thresholds(samples.cuda(), labels.cuda(), device='cuda'),
+        )

From 0d864695b31b8ac13f6a3c416d3f7932fe0a869c Mon Sep 17 00:00:00 2001
From: narugo1992 <narugo1992@users.noreply.github.com>
Date: Mon, 10 Aug 2026 11:43:39 +0000
Subject: [PATCH 2/4] dev(narugo): validate histogram inputs found during
 self-review

Reject histograms whose bin count does not match num_thresholds. The mismatch
would not have raised on its own -- a 50-bin argmax indexes a 100-point grid
without going out of bounds -- so it would have silently reported thresholds
taken from the wrong grid. Also reject mismatched histogram pairs, non-2D
batches, and an empty threshold grid, and make the float64 comparison path
explicit so float64 inputs keep the precision the original numpy comparison had.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 animetimm/multilabel/metrics.py | 31 ++++++++++++++++++++-------
 requirements-test.txt           |  1 -
 test/multilabel/test_metrics.py | 37 +++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 9 deletions(-)

diff --git a/animetimm/multilabel/metrics.py b/animetimm/multilabel/metrics.py
index d0d5ac9..6b6ed62 100644
--- a/animetimm/multilabel/metrics.py
+++ b/animetimm/multilabel/metrics.py
@@ -83,6 +83,8 @@ def recall(tp, fp, tn, fn, mean: bool = True):
 
 def threshold_grid(num_thresholds: int = 100) -> np.ndarray:
     """Candidate thresholds, ``num_thresholds`` points evenly spread over ``(0, 1]``."""
+    if num_thresholds < 1:
+        raise ValueError(f'num_thresholds must be at least 1, got {num_thresholds!r}.')
     return np.linspace(1.0 / num_thresholds, 1, num_thresholds)
 
 
@@ -152,7 +154,7 @@ def _allocate(self, device):
     def _bounds_for(self, sample):
         """Bin edges, built from the same np.linspace the offline path uses so both
         routes bucketize against bit-identical values."""
-        dtype = self.score_dtype or _bound_dtype(sample.dtype)
+        dtype = self.score_dtype if self.score_dtype is not None else _bound_dtype(sample.dtype)
         if self._bounds is None or self._bounds.dtype != dtype:
             self._bounds = torch.as_tensor(threshold_grid(self.num_thresholds),
                                            dtype=dtype, device=self._device)
@@ -167,6 +169,8 @@ def _ones_like_rows(self, rows: int):
     def update(self, sample, labels):
         """Fold one batch in. ``sample`` holds sigmoid scores, ``labels`` the raw labels,
         both shaped ``(batch_size, tag_num)``."""
+        if sample.ndim != 2:
+            raise ValueError(f'Expected a (batch_size, tag_num) batch, got {sample.ndim} dimensions.')
         if sample.shape != labels.shape:
             raise ValueError(f'Sample shape {tuple(sample.shape)} does not match '
                              f'label shape {tuple(labels.shape)}.')
@@ -186,7 +190,7 @@ def finalize(self, accelerator=None) -> Histograms:
         """Reduce across processes when running distributed, and hand back the
         ``(tag_num, num_thresholds + 1)`` histograms the searches consume."""
         if self.hist_all is None:
-            self._allocate(self._device or torch.device('cpu'))
+            self._allocate(self._device if self._device is not None else torch.device('cpu'))
 
         hist_all, hist_pos = self.hist_all, self.hist_pos
         if accelerator is not None and accelerator.num_processes > 1:
@@ -204,6 +208,8 @@ def build_threshold_histograms(all_sample, all_labels, num_thresholds: int = 100
 
     Tags are processed in chunks so peak memory stays bounded regardless of tag count.
     """
+    if all_sample.ndim != 2:
+        raise ValueError(f'Expected a (sample_num, tag_num) matrix, got {all_sample.ndim} dimensions.')
     if all_sample.shape != all_labels.shape:
         raise ValueError(f'Sample shape {tuple(all_sample.shape)} does not match '
                          f'label shape {tuple(all_labels.shape)}.')
@@ -281,12 +287,21 @@ def _pick(f1s, pres, recs, ths):
 
 
 def _resolve_histograms(all_sample, all_labels, num_thresholds, histograms, device):
-    if histograms is not None:
-        return histograms
-    if all_sample is None or all_labels is None:
-        raise ValueError('Either histograms, or both all_sample and all_labels, must be given.')
-    return build_threshold_histograms(all_sample, all_labels, num_thresholds=num_thresholds,
-                                      device=device)
+    if histograms is None:
+        if all_sample is None or all_labels is None:
+            raise ValueError('Either histograms, or both all_sample and all_labels, must be given.')
+        return build_threshold_histograms(all_sample, all_labels, num_thresholds=num_thresholds,
+                                          device=device)
+
+    hist_all, hist_pos = histograms
+    # a grid of the wrong size would still index without raising, and would silently
+    # report thresholds taken from the wrong grid
+    if hist_all.shape[-1] != num_thresholds + 1:
+        raise ValueError(f'Histograms have {hist_all.shape[-1]} bins, which corresponds to '
+                         f'{hist_all.shape[-1] - 1} thresholds, but num_thresholds is {num_thresholds}.')
+    if hist_all.shape != hist_pos.shape:
+        raise ValueError(f'Histogram shapes disagree: {hist_all.shape} versus {hist_pos.shape}.')
+    return histograms
 
 
 def compute_optimal_thresholds(all_sample=None, all_labels=None, alpha: float = 1.0, num_thresholds: int = 100,
diff --git a/requirements-test.txt b/requirements-test.txt
index 447807c..d5defa9 100644
--- a/requirements-test.txt
+++ b/requirements-test.txt
@@ -1,3 +1,2 @@
 pytest~=7.4
 pytest-cov~=4.1
-pytest-xdist~=3.5
diff --git a/test/multilabel/test_metrics.py b/test/multilabel/test_metrics.py
index 50c211c..41b331b 100644
--- a/test/multilabel/test_metrics.py
+++ b/test/multilabel/test_metrics.py
@@ -170,6 +170,11 @@ def test_endpoints(self, num_thresholds):
         assert grid[-1] == pytest.approx(1.0)
         assert grid[0] == pytest.approx(1.0 / num_thresholds)
 
+    @pytest.mark.parametrize('num_thresholds', [0, -1])
+    def test_rejects_empty_grid(self, num_thresholds):
+        with pytest.raises(ValueError, match='at least 1'):
+            threshold_grid(num_thresholds)
+
 
 @pytest.mark.unittest
 class TestPerTagSearch:
@@ -288,6 +293,32 @@ def test_rejects_shape_mismatch(self):
         with pytest.raises(ValueError, match='does not match'):
             build_threshold_histograms(samples, labels[:, :2])
 
+    def test_rejects_non_2d_input(self):
+        samples, labels = make_case(sample_num=64, tag_num=4)
+        with pytest.raises(ValueError, match='got 1 dimensions'):
+            build_threshold_histograms(samples[0], labels[0])
+
+    def test_rejects_histograms_built_for_a_different_grid(self):
+        """Indexing a 100-point grid with a 50-bin argmax would not raise on its own, it
+        would just report thresholds off the wrong grid."""
+        samples, labels = make_case(sample_num=128, tag_num=6, seed=25)
+        coarse = build_threshold_histograms(samples, labels, num_thresholds=50)
+        with pytest.raises(ValueError, match='num_thresholds is 100'):
+            compute_optimal_thresholds(histograms=coarse)
+        with pytest.raises(ValueError, match='num_thresholds is 100'):
+            compute_optimal_thresholds_by_categories(df_tags=category_frame(6), histograms=coarse)
+        # ... and it works when the caller says which grid the histograms belong to
+        assert_all_identical(
+            reference_optimal_thresholds(samples, labels, num_thresholds=50),
+            compute_optimal_thresholds(histograms=coarse, num_thresholds=50),
+        )
+
+    def test_rejects_mismatched_histogram_pair(self):
+        samples, labels = make_case(sample_num=128, tag_num=6, seed=26)
+        hist_all, hist_pos = build_threshold_histograms(samples, labels)
+        with pytest.raises(ValueError, match='shapes disagree'):
+            compute_optimal_thresholds(histograms=(hist_all, hist_pos[:3]))
+
 
 @pytest.mark.unittest
 class TestStreamingHistogram:
@@ -406,6 +437,12 @@ def test_rejects_wrong_tag_count(self):
         with pytest.raises(ValueError, match='Expected 5 tags'):
             accumulator.update(samples, labels)
 
+    def test_rejects_non_2d_batch(self):
+        samples, labels = make_case(sample_num=32, tag_num=4)
+        accumulator = StreamingThresholdHistogram(4)
+        with pytest.raises(ValueError, match='got 1 dimensions'):
+            accumulator.update(samples[0], labels[0])
+
 
 @pytest.mark.unittest
 class TestScoreDtypes:

From 5a3024cc1766b321efbf775823d38fe1ea52e9bc Mon Sep 17 00:00:00 2001
From: narugo1992 <narugo1992@users.noreply.github.com>
Date: Mon, 10 Aug 2026 11:52:33 +0000
Subject: [PATCH 3/4] ci(narugo): run the multilabel metrics unit tests

Runs the threshold-search suite on CPU-only wheels, which is also what proves the
CPU-only path keeps working: the CUDA parity tests skip and the remaining 58 must
still pass.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 .github/workflows/unittest.yml | 45 ++++++++++++++++++++++++++++++++++
 1 file changed, 45 insertions(+)
 create mode 100644 .github/workflows/unittest.yml

diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
new file mode 100644
index 0000000..3095184
--- /dev/null
+++ b/.github/workflows/unittest.yml
@@ -0,0 +1,45 @@
+#file: noinspection YAMLSchemaValidation
+name: Unittest
+
+on:
+  push:
+    paths:
+      - 'animetimm/**'
+      - 'test/**'
+      - 'requirements*.txt'
+      - '.github/workflows/unittest.yml'
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  unittest:
+    name: Unittest
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os:
+          - 'ubuntu-latest'
+        python-version:
+          - '3.10'
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 1
+      - name: Set up python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        shell: bash
+        run: |
+          python -m pip install --upgrade pip
+          pip install --index-url https://download.pytorch.org/whl/cpu 'torch>=2.6.0,<2.9.0'
+          pip install numpy pandas 'accelerate>=1.4.0'
+          pip install -r requirements-test.txt
+      - name: Run unittest
+        shell: bash
+        run: |
+          python -m pytest test/multilabel/test_metrics.py -sv -m unittest

From d253e76f0f7d7093bb7cee3f966aabdced1a2d7e Mon Sep 17 00:00:00 2001
From: narugo1992 <narugo1992@users.noreply.github.com>
Date: Mon, 10 Aug 2026 12:47:04 +0000
Subject: [PATCH 4/4] dev(narugo): compare scores against unrounded thresholds,
 and test the maths without a GPU

Two things this fixes, both found by asking whether a CPU-only run can actually
decide correctness.

Comparison dtype. The thresholds are float64 and the original scan wrote
`sample >= th` in numpy, which promotes the float32 score to float64. Bucketizing
against float32 bin edges rounds the threshold instead, which flips any score
within one float32 ulp of one: float32(0.03) is 0.0299999993, below the real 0.03
threshold, yet not below its float32 rounding. Sweeping four float32 steps either
side of all 100 thresholds, float32 edges disagree with the original on 50 of 800
scores while float64 edges disagree on none. Real sigmoid outputs never land there
-- zero disagreements across 2.5e8 measured scores, which is why the full-scale
comparison was bit-identical either way -- but the promotion removes the
discrepancy outright and costs little: a-gpu 7.10s -> 8.92s, b-gpu 0.117 ->
0.125 ms/batch, b-cpu unchanged.

Test coverage without a GPU. mcc/f1score/precision/recall had no tests at all
despite feeding test_metrics.json, and were verified only by agreement with the
old code, which cannot catch an error both share. They now have hand-worked
expected values, with MCC additionally checked against the textbook formula, plus
worked examples for the histogram, the threshold search and the tie-walk. The
device-transfer branch in update() was unreachable on CPU, so it is gone: .to() of
a tensor already on the target device is a no-op. Coverage of metrics.py on a
CPU-only run is now 100%, and CI enforces it, so a line only a CUDA machine can
reach fails the build rather than quietly going untested.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 .github/workflows/unittest.yml  |   9 +-
 animetimm/multilabel/metrics.py |  24 ++-
 test/multilabel/test_metrics.py | 258 +++++++++++++++++++++++++++++---
 3 files changed, 265 insertions(+), 26 deletions(-)

diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
index 3095184..cff7bec 100644
--- a/.github/workflows/unittest.yml
+++ b/.github/workflows/unittest.yml
@@ -42,4 +42,11 @@ jobs:
       - name: Run unittest
         shell: bash
         run: |
-          python -m pytest test/multilabel/test_metrics.py -sv -m unittest
+          # --cov-fail-under=100 is deliberate: this runner has no GPU, so the gate is
+          # what guarantees the threshold maths stays fully decidable without one. A new
+          # line that only a CUDA machine can reach will fail here rather than quietly
+          # becoming untested.
+          python -m pytest test/multilabel/test_metrics.py -sv -m unittest \
+            --cov=animetimm.multilabel.metrics \
+            --cov-report=term-missing \
+            --cov-fail-under=100
diff --git a/animetimm/multilabel/metrics.py b/animetimm/multilabel/metrics.py
index 6b6ed62..8d819f4 100644
--- a/animetimm/multilabel/metrics.py
+++ b/animetimm/multilabel/metrics.py
@@ -95,13 +95,19 @@ def _binarize(labels):
 
 
 def _bound_dtype(sample_dtype: torch.dtype) -> torch.dtype:
-    """Precision to compare scores against thresholds in.
-
-    float64 inputs keep full precision, matching what the original numpy comparison did;
-    everything else (float32 in practice, float16 in principle) compares in float32,
-    since a 0.01 grid is not representable well enough in half.
+    """Precision to compare scores against thresholds in: always float64.
+
+    The thresholds come from ``np.linspace``, so they are float64 values, and the original
+    scan compared with ``sample >= th`` in numpy, which promotes the float32 scores to
+    float64 before comparing. Bucketizing in float32 instead would round the *thresholds*
+    and flip scores that sit within one float32 ulp of one -- float32(0.03) is
+    0.0299999993, which the original counts as below the 0.03 threshold but a float32
+    comparison counts as at it. Real sigmoid outputs essentially never land there (zero
+    disagreements over 2.5e8 measured scores), but the promotion costs a temporary in the
+    batch being bucketized and removes the discrepancy entirely.
     """
-    return torch.float64 if sample_dtype == torch.float64 else torch.float32
+    _ = sample_dtype
+    return torch.float64
 
 
 def _bin_indices(sample, bounds):
@@ -178,8 +184,10 @@ def update(self, sample, labels):
             raise ValueError(f'Expected {self.tag_num} tags, got {sample.shape[-1]}.')
         if self.hist_all is None:
             self._allocate(sample.device)
-        elif sample.device != self._device:
-            sample, labels = sample.to(self._device), labels.to(self._device)
+        # unconditional rather than guarded: .to() of a tensor already on the target
+        # device returns it unchanged, and a guard here would be a branch that a
+        # CPU-only test run could never reach
+        sample, labels = sample.to(self._device), labels.to(self._device)
 
         bins = _bin_indices(sample, self._bounds_for(sample))
         self.hist_all.scatter_add_(0, bins, self._ones_like_rows(bins.shape[0]))
diff --git a/test/multilabel/test_metrics.py b/test/multilabel/test_metrics.py
index 41b331b..43e224d 100644
--- a/test/multilabel/test_metrics.py
+++ b/test/multilabel/test_metrics.py
@@ -15,6 +15,10 @@
     build_threshold_histograms,
     compute_optimal_thresholds,
     compute_optimal_thresholds_by_categories,
+    f1score,
+    mcc,
+    precision,
+    recall,
     threshold_grid,
 )
 
@@ -156,6 +160,160 @@ def assert_all_identical(expected, actual, names=('threshold', 'f1', 'precision'
 
 # ------------------------------------------------------------------------- tests
 
+@pytest.mark.unittest
+class TestHandWorkedExamples:
+    """Absolute correctness, not just agreement with the previous implementation.
+
+    Every expected value below is derived by hand in the comments, so these tests would
+    catch an error that the old code and the new code happen to share. They need no GPU,
+    which is the point: correctness is decidable on a CPU-only machine.
+    """
+
+    # 5 samples, 1 tag, K=10 -> thresholds 0.1, 0.2, ... 1.0
+    SCORES = [0.05, 0.15, 0.35, 0.55, 0.95]
+    LABELS = [0, 0, 1, 1, 1]
+
+    def _one_tag(self):
+        return (torch.tensor([[s] for s in self.SCORES], dtype=torch.float32),
+                torch.tensor([[float(v)] for v in self.LABELS], dtype=torch.float32))
+
+    def test_bin_indices_are_the_count_of_thresholds_at_or_below_the_score(self):
+        # bin(s) = #{k : threshold[k] <= s}
+        #   0.05 -> 0 thresholds   0.15 -> 1 (0.1)     0.35 -> 3 (0.1,0.2,0.3)
+        #   0.55 -> 5 (..0.5)      0.95 -> 9 (..0.9)
+        hist_all, hist_pos = build_threshold_histograms(*self._one_tag(), num_thresholds=10)
+        assert {i: int(v) for i, v in enumerate(hist_all[0]) if v} == {0: 1, 1: 1, 3: 1, 5: 1, 9: 1}
+        # only the last three samples are positive
+        assert {i: int(v) for i, v in enumerate(hist_pos[0]) if v} == {3: 1, 5: 1, 9: 1}
+
+    def test_optimal_threshold_matches_the_hand_derivation(self):
+        # th=0.1 -> tp=3 fp=1 fn=0 -> p=3/4 r=1    f1=0.857142...
+        # th=0.2 -> tp=3 fp=0 fn=0 -> p=1   r=1    f1=1.0        <- joint best
+        # th=0.3 -> tp=3 fp=0 fn=0 -> p=1   r=1    f1=1.0        <- joint best
+        # th=0.4 -> tp=2 fp=0 fn=1 -> p=1   r=2/3  f1=0.8
+        # ... decreasing from there, and 0 at th=1.0
+        # the run of identical (f1, p, r) spans 0.2 and 0.3, so the midpoint is 0.25
+        th, f1, p, r = compute_optimal_thresholds(*self._one_tag(), num_thresholds=10)
+        assert th[0] == pytest.approx(0.25)
+        assert f1[0] == pytest.approx(1.0)
+        assert p[0] == pytest.approx(1.0)
+        assert r[0] == pytest.approx(1.0)
+
+    def test_tie_run_spanning_many_thresholds_takes_the_midpoint(self):
+        # one positive sample scoring 0.95: every threshold from 0.1 to 0.9 predicts it
+        # correctly (f1 = 1.0), and 1.0 misses it (f1 = 0). The run is 0.1 .. 0.9, so the
+        # midpoint is (0.1 + 0.9) / 2 = 0.5
+        samples = torch.tensor([[0.95]], dtype=torch.float32)
+        labels = torch.tensor([[1.0]], dtype=torch.float32)
+        th, f1, _, _ = compute_optimal_thresholds(samples, labels, num_thresholds=10)
+        assert th[0] == pytest.approx(0.5)
+        assert f1[0] == pytest.approx(1.0)
+
+    def test_category_result_is_the_micro_average_over_its_tags(self):
+        # tag0: scores [0.15, 0.35, 0.55], labels [0, 1, 1]
+        # tag1: scores [0.25, 0.45, 0.05], labels [1, 1, 0]
+        # th=0.1 -> tp=4 fp=1 fn=0 -> p=0.8 r=1    f1=0.888...
+        # th=0.2 -> tp=4 fp=0 fn=0 -> p=1   r=1    f1=1.0     <- unique best
+        # th=0.3 -> tp=3 fp=0 fn=1 -> p=1   r=0.75 f1=0.857...
+        samples = torch.tensor([[0.15, 0.25], [0.35, 0.45], [0.55, 0.05]], dtype=torch.float32)
+        labels = torch.tensor([[0.0, 1.0], [1.0, 1.0], [1.0, 0.0]], dtype=torch.float32)
+        df_tags = pd.DataFrame({'category': [7, 7]})
+        th, f1, p, r = compute_optimal_thresholds_by_categories(samples, labels, df_tags,
+                                                                num_thresholds=10)
+        assert th[7] == pytest.approx(0.2)
+        assert f1[7] == pytest.approx(1.0)
+        assert p[7] == pytest.approx(1.0)
+        assert r[7] == pytest.approx(1.0)
+
+    def test_perfect_and_inverted_separation(self):
+        # tag 0 is perfectly separable at 0.5; tag 1 has scores exactly inverted, so no
+        # threshold does better than predicting everything positive
+        samples = torch.tensor([[0.9, 0.1], [0.8, 0.2], [0.2, 0.8], [0.1, 0.9]], dtype=torch.float32)
+        labels = torch.tensor([[1.0, 1.0], [1.0, 1.0], [0.0, 0.0], [0.0, 0.0]], dtype=torch.float32)
+        _, f1, p, r = compute_optimal_thresholds(samples, labels, num_thresholds=10)
+        assert f1[0] == pytest.approx(1.0) and p[0] == pytest.approx(1.0) and r[0] == pytest.approx(1.0)
+        # inverted: best is threshold 0.1, catching both positives plus both negatives,
+        # so tp=2 fp=2 fn=0 -> p=0.5, r=1, f1=2/3
+        assert f1[1] == pytest.approx(2 / 3)
+        assert p[1] == pytest.approx(0.5)
+        assert r[1] == pytest.approx(1.0)
+
+
+@pytest.mark.unittest
+class TestConfusionMetrics:
+    """`mcc`, `f1score`, `precision` and `recall` feed test_metrics.json directly.
+
+    Checked against values worked out by hand, and MCC additionally against the textbook
+    formula, since the implementation uses a rescaled form that is not obviously the same.
+    """
+
+    @staticmethod
+    def counts(tp, fp, tn, fn):
+        return tuple(torch.tensor([float(v)]) for v in (tp, fp, tn, fn))
+
+    def test_precision_recall_f1_by_hand(self):
+        # tp=3 fp=3 tn=4 fn=1
+        args = self.counts(3, 3, 4, 1)
+        assert precision(*args).item() == pytest.approx(3 / 6)
+        assert recall(*args).item() == pytest.approx(3 / 4)
+        assert f1score(*args).item() == pytest.approx(2 * 3 / (2 * 3 + 1 + 3), rel=1e-6)
+
+    def test_mcc_equals_the_textbook_formula(self):
+        for tp, fp, tn, fn in [(3, 3, 4, 1), (10, 2, 50, 8), (1, 0, 99, 0), (25, 25, 25, 25)]:
+            expected = ((tp * tn - fp * fn)
+                        / np.sqrt(float((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))))
+            assert mcc(*self.counts(tp, fp, tn, fn)).item() == pytest.approx(expected, rel=1e-5), \
+                f'tp={tp} fp={fp} tn={tn} fn={fn}'
+
+    def test_perfect_and_worst_case(self):
+        assert mcc(*self.counts(50, 0, 50, 0)).item() == pytest.approx(1.0, rel=1e-5)
+        assert f1score(*self.counts(50, 0, 50, 0)).item() == pytest.approx(1.0, rel=1e-6)
+        # every prediction inverted
+        assert mcc(*self.counts(0, 50, 0, 50)).item() == pytest.approx(-1.0, rel=1e-5)
+
+    def test_alpha_is_beta_squared_not_beta(self):
+        """Locks a quirk worth knowing: `f1score`'s alpha is the *squared* beta of the
+        usual F-beta, while `compute_optimal_thresholds` squares its own alpha. Both are
+        called with alpha=1.0 in this repository, where the two conventions coincide."""
+        args = self.counts(3, 3, 4, 1)   # p = 0.5, r = 0.75
+        p_, r_ = 0.5, 0.75
+        f_beta_sq_4 = (1 + 4) * p_ * r_ / (4 * p_ + r_)
+        f_beta_4 = (1 + 16) * p_ * r_ / (16 * p_ + r_)
+        assert f1score(*args, alpha=4).item() == pytest.approx(f_beta_sq_4, rel=1e-6)
+        assert f1score(*args, alpha=4).item() != pytest.approx(f_beta_4, rel=1e-6)
+
+    @pytest.mark.parametrize('tp,fp,tn,fn', [
+        (0, 0, 10, 0),    # nothing predicted positive, nothing actually positive
+        (0, 10, 0, 0),    # everything predicted positive, nothing actually positive
+        (0, 0, 0, 10),    # nothing predicted positive, everything actually positive
+        (10, 0, 0, 0),    # everything positive and every prediction right
+    ])
+    def test_degenerate_counts_do_not_produce_nan(self, tp, fp, tn, fn):
+        for fn_ in (mcc, f1score, precision, recall):
+            value = fn_(*self.counts(tp, fp, tn, fn)).item()
+            assert not np.isnan(value), f'{fn_.__name__} returned NaN for {(tp, fp, tn, fn)}'
+
+    def test_empty_confusion_matrix_yields_nan_mcc(self):
+        """Documents pre-existing behaviour rather than endorsing it: with no samples at
+        all, `mcc` divides 0 by 0 while the other three return finite values. This is
+        unreachable in `test.py` — micro counts total over every sample and tag, macro
+        counts total over every sample — so it is recorded, not fixed, here."""
+        empty = self.counts(0, 0, 0, 0)
+        assert np.isnan(mcc(*empty).item())
+        for fn_ in (f1score, precision, recall):
+            assert not np.isnan(fn_(*empty).item())
+
+    def test_mean_false_returns_per_entry_values(self):
+        tp = torch.tensor([3.0, 10.0])
+        fp = torch.tensor([3.0, 2.0])
+        tn = torch.tensor([4.0, 50.0])
+        fn = torch.tensor([1.0, 8.0])
+        per_entry = precision(tp, fp, tn, fn, mean=False)
+        assert per_entry.shape == (2,)
+        np.testing.assert_allclose(per_entry.numpy(), [3 / 6, 10 / 12], rtol=1e-6)
+        assert precision(tp, fp, tn, fn).item() == pytest.approx((3 / 6 + 10 / 12) / 2, rel=1e-6)
+
+
 @pytest.mark.unittest
 class TestThresholdGrid:
     def test_matches_documented_grid(self):
@@ -463,24 +621,90 @@ def test_streaming_matches_offline(self, dtype):
         np.testing.assert_array_equal(build_threshold_histograms(samples, labels)[0],
                                       accumulator.finalize()[0])
 
-    def test_float64_input_compares_in_float64(self):
-        """A score just below a threshold in float64, but indistinguishable from it once
-        rounded to float32, must be judged the way the original float64 numpy comparison
-        judged it: below."""
+    @pytest.mark.parametrize('dtype', [torch.float16, torch.bfloat16])
+    def test_half_precision_scores_are_promoted(self, dtype):
+        """Mixed-precision inference can hand back half-precision scores. A 0.01 grid is
+        nowhere near representable in half, so the comparison is promoted; numpy cannot
+        even hold bfloat16, so the old code would have failed outright here."""
+        scores = [0.05, 0.15, 0.35, 0.55, 0.95]
+        labels = [0.0, 0.0, 1.0, 1.0, 1.0]
+        half = torch.tensor([[s] for s in scores], dtype=dtype)
+        half_labels = torch.tensor([[v] for v in labels], dtype=dtype)
+
+        hist_all, hist_pos = build_threshold_histograms(half, half_labels, num_thresholds=10)
+        # same bins as the float32 hand-worked example: rounding to half moves each score
+        # by less than the distance to its neighbouring threshold
+        assert {i: int(v) for i, v in enumerate(hist_all[0]) if v} == {0: 1, 1: 1, 3: 1, 5: 1, 9: 1}
+        assert {i: int(v) for i, v in enumerate(hist_pos[0]) if v} == {3: 1, 5: 1, 9: 1}
+
+        th, f1, _, _ = compute_optimal_thresholds(half, half_labels, num_thresholds=10)
+        assert th[0] == pytest.approx(0.25)
+        assert f1[0] == pytest.approx(1.0)
+
+        # and the streaming path agrees with the offline one on the same input
+        accumulator = StreamingThresholdHistogram(1, num_thresholds=10)
+        accumulator.update(half, half_labels)
+        np.testing.assert_array_equal(hist_all, accumulator.finalize()[0])
+
+    def test_half_promotion_changes_a_boundary_score(self):
+        """Pins the promotion down with a score the precisions disagree about.
+
+        float16 rounds 0.1 *down* to 0.099975..., below the true 0.1 threshold, so a score
+        sitting exactly on the half representation of the first threshold falls below it
+        once promoted, and lands in bin 0. Comparing in half would put it in bin 1.
+        """
+        boundary = float(np.float16(0.1))
+        assert boundary < 0.1
+
+        scores = torch.tensor([[boundary]], dtype=torch.float16)
+        labels = torch.ones((1, 1), dtype=torch.float16)
+        hist_all, _ = build_threshold_histograms(scores, labels, num_thresholds=10)
+        assert int(np.argmax(hist_all[0])) == 0
+
+        accumulator = StreamingThresholdHistogram(1, num_thresholds=10)
+        accumulator.update(scores, labels)
+        assert int(np.argmax(accumulator.finalize()[0][0])) == 0
+
+    def test_float32_scores_are_judged_against_unrounded_thresholds(self):
+        """The thresholds are float64 and the original compared `sample >= th` in numpy,
+        which promotes the float32 score rather than rounding the threshold. Bucketizing
+        in float32 would round the threshold instead and flip scores within one ulp of it:
+        float32(0.03) is 0.0299999993, below the real 0.03, yet not below float32(0.03).
+        """
+        grid = threshold_grid(100)
+        score = np.float32(grid[2])                   # the float32 image of the 0.03 edge
+        assert float(score) < grid[2]                 # genuinely below the real threshold
+        assert score >= np.float32(grid[2])           # but not below its float32 rounding
+
+        samples = torch.tensor([[float(score)]], dtype=torch.float32)
+        hist_all, _ = build_threshold_histograms(samples, torch.ones_like(samples))
+        assert int(np.argmax(hist_all[0])) == 2       # 0.01 and 0.02 only, not 0.03
+
+    def test_matches_original_numpy_semantics_on_every_threshold_boundary(self):
+        """Sweep four float32 steps either side of every threshold and require the bins to
+        agree with `score.astype(float64) >= thresholds`, which is what the original scan
+        computed."""
         grid = threshold_grid(100)
-        edge = grid[6] - 1e-12                       # just under the 0.07 threshold
-        assert edge < grid[6]                        # ... in float64
-        assert np.float32(edge) >= np.float32(grid[6])  # ... but not in float32
-
-        labels = torch.ones((1, 1), dtype=torch.float64)
-        hist_all, _ = build_threshold_histograms(torch.tensor([[edge]], dtype=torch.float64), labels)
-        # the bin counts thresholds <= score, so 0.07 must not be counted
-        assert int(np.argmax(hist_all[0])) == 6
-
-        # the same value handed in as float32 genuinely is >= the threshold
-        hist32, _ = build_threshold_histograms(torch.tensor([[edge]], dtype=torch.float32),
-                                               labels.to(torch.float32))
-        assert int(np.argmax(hist32[0])) == 7
+        edge = []
+        for th in grid:
+            v = np.float32(th)
+            for _ in range(4):
+                edge.append(v)
+                v = np.nextafter(v, np.float32(0))
+            v = np.float32(th)
+            for _ in range(4):
+                v = np.nextafter(v, np.float32(2))
+                edge.append(v)
+        edge = np.array(edge, dtype=np.float32)
+        expected = (edge[:, None].astype(np.float64) >= grid[None, :]).sum(axis=1)
+
+        # one score per tag, so each row's histogram has a single entry and its argmax is
+        # exactly that score's bin
+        samples = torch.from_numpy(edge.reshape(1, -1))
+        hist_all, _ = build_threshold_histograms(samples, torch.ones_like(samples))
+        assert hist_all.shape == (edge.size, 101)
+        np.testing.assert_array_equal(hist_all.sum(axis=1), np.ones(edge.size))
+        np.testing.assert_array_equal(expected, hist_all.argmax(axis=1))
 
 
 @pytest.mark.unittest