Repository: lenskit/lkpy

Issue:
Use PyTorch to implement user-user k-NN
Advance #374 by using PyTorch to implement user-user k-NN.

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/lenskit/lkpy/pull/406?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lenskit) Report
Attention: Patch coverage is `39.34426%` with `37 lines` in your changes are missing coverage. Please review.
> Project coverage is 86.79%. Comparing base [(`39ff9f5`)](https://app.codecov.io/gh/lenskit/lkpy/commit/39ff9f549ac63f18da0902ca6d7eb3719c71405d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lenskit) to head [(`9c4726a`)](https://app.codecov.io/gh/lenskit/lkpy/commit/9c4726a562fb0af86baee180e905080ec695c0f2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lenskit).

| [Files](https://app.codecov.io/gh/lenskit/lkpy/pull/406?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lenskit) | Patch % | Lines |
|---|---|---|
| [lenskit/algorithms/knn/user.py](https://app.codecov.io/gh/lenskit/lkpy/pull/406?src=pr&el=tree&filepath=lenskit%2Falgorithms%2Fknn%2Fuser.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lenskit#diff-bGVuc2tpdC9hbGdvcml0aG1zL2tubi91c2VyLnB5) | 13.95% | [37 Missing :warning: ](https://app.codecov.io/gh/lenskit/lkpy/pull/406?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lenskit) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #406      +/-   ##
==========================================
+ Coverage   86.27%   86.79%   +0.52%     
==========================================
  Files          48       48              
  Lines        3212     3203       -9     
==========================================
+ Hits         2771     2780       +9     
+ Misses        441      423      -18     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/lenskit/lkpy/pull/406?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lenskit).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lenskit).


Patch:
diff --git a/envs/lenskit-py3.10-ci.yaml b/envs/lenskit-py3.10-ci.yaml
index eac2068d..1a7cd0ae 100644
--- a/envs/lenskit-py3.10-ci.yaml
+++ b/envs/lenskit-py3.10-ci.yaml
@@ -8,8 +8,8 @@
 # Instead edit the corresponding pyproject.toml file.
 #
 channels:
-  - conda-forge
   - pytorch
+  - conda-forge
   - nodefaults
 dependencies:
   - python=3.10
@@ -30,9 +30,10 @@ dependencies:
   - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus==1.*
+  - pytest-repeat>=0.9
   - pytest==7.*
   - python-build==1
-  - pytorch==2.*
+  - pytorch<3,>=2.1
   - ruff>=0.2
   - scikit-learn>=1.1
   - scipy>=1.9.0
diff --git a/envs/lenskit-py3.10-dev.yaml b/envs/lenskit-py3.10-dev.yaml
index e9573b00..b924bdad 100644
--- a/envs/lenskit-py3.10-dev.yaml
+++ b/envs/lenskit-py3.10-dev.yaml
@@ -8,8 +8,8 @@
 # Instead edit the corresponding pyproject.toml file.
 #
 channels:
-  - conda-forge
   - pytorch
+  - conda-forge
   - nodefaults
 dependencies:
   - python=3.10
@@ -35,9 +35,10 @@ dependencies:
   - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus==1.*
+  - pytest-repeat>=0.9
   - pytest==7.*
   - python-build==1
-  - pytorch==2.*
+  - pytorch<3,>=2.1
   - ruff>=0.2
   - scikit-learn>=1.1
   - scipy>=1.9.0
diff --git a/envs/lenskit-py3.10-doc.yaml b/envs/lenskit-py3.10-doc.yaml
index 8a6f4dfa..dba22f7f 100644
--- a/envs/lenskit-py3.10-doc.yaml
+++ b/envs/lenskit-py3.10-doc.yaml
@@ -8,8 +8,8 @@
 # Instead edit the corresponding pyproject.toml file.
 #
 channels:
-  - conda-forge
   - pytorch
+  - conda-forge
   - nodefaults
 dependencies:
   - python=3.10
@@ -20,7 +20,7 @@ dependencies:
   - numba<0.59,>=0.56
   - numpy>=1.23
   - pandas<3,>=1.5
-  - pytorch==2.*
+  - pytorch<3,>=2.1
   - scipy>=1.9.0
   - sphinx>=4.2
   - sphinx_rtd_theme>=0.5
diff --git a/envs/lenskit-py3.11-ci.yaml b/envs/lenskit-py3.11-ci.yaml
index 51035850..f3577218 100644
--- a/envs/lenskit-py3.11-ci.yaml
+++ b/envs/lenskit-py3.11-ci.yaml
@@ -8,8 +8,8 @@
 # Instead edit the corresponding pyproject.toml file.
 #
 channels:
-  - conda-forge
   - pytorch
+  - conda-forge
   - nodefaults
 dependencies:
   - python=3.11
@@ -30,9 +30,10 @@ dependencies:
   - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus==1.*
+  - pytest-repeat>=0.9
   - pytest==7.*
   - python-build==1
-  - pytorch==2.*
+  - pytorch<3,>=2.1
   - ruff>=0.2
   - scikit-learn>=1.1
   - scipy>=1.9.0
diff --git a/envs/lenskit-py3.11-dev.yaml b/envs/lenskit-py3.11-dev.yaml
index 0989c30c..87eccc4d 100644
--- a/envs/lenskit-py3.11-dev.yaml
+++ b/envs/lenskit-py3.11-dev.yaml
@@ -8,8 +8,8 @@
 # Instead edit the corresponding pyproject.toml file.
 #
 channels:
-  - conda-forge
   - pytorch
+  - conda-forge
   - nodefaults
 dependencies:
   - python=3.11
@@ -35,9 +35,10 @@ dependencies:
   - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus==1.*
+  - pytest-repeat>=0.9
   - pytest==7.*
   - python-build==1
-  - pytorch==2.*
+  - pytorch<3,>=2.1
   - ruff>=0.2
   - scikit-learn>=1.1
   - scipy>=1.9.0
diff --git a/envs/lenskit-py3.11-doc.yaml b/envs/lenskit-py3.11-doc.yaml
index dd78379d..443c85df 100644
--- a/envs/lenskit-py3.11-doc.yaml
+++ b/envs/lenskit-py3.11-doc.yaml
@@ -8,8 +8,8 @@
 # Instead edit the corresponding pyproject.toml file.
 #
 channels:
-  - conda-forge
   - pytorch
+  - conda-forge
   - nodefaults
 dependencies:
   - python=3.11
@@ -20,7 +20,7 @@ dependencies:
   - numba<0.59,>=0.56
   - numpy>=1.23
   - pandas<3,>=1.5
-  - pytorch==2.*
+  - pytorch<3,>=2.1
   - scipy>=1.9.0
   - sphinx>=4.2
   - sphinx_rtd_theme>=0.5
diff --git a/lenskit/algorithms/knn/user.py b/lenskit/algorithms/knn/user.py
index 35d5bed5..0442a506 100644
--- a/lenskit/algorithms/knn/user.py
+++ b/lenskit/algorithms/knn/user.py
@@ -13,17 +13,15 @@ from __future__ import annotations
 
 import logging
 import warnings
-from typing import NamedTuple, Optional
 
 import numpy as np
 import pandas as pd
 import torch
-from numba import njit
+from typing_extensions import Any, NamedTuple, Optional, Self
 
 from lenskit import DataWarning, util
 from lenskit.data import FeedbackType, sparse_ratings
-from lenskit.data.matrix import normalize_sparse_rows, sparse_row_stats
-from lenskit.util.accum import kvp_minheap_insert
+from lenskit.data.matrix import normalize_sparse_rows, safe_spmv
 
 from .. import Predictor
 
@@ -81,8 +79,10 @@ class UserUser(Predictor):
     aggregate: str
     use_ratings: bool
 
-    user_index_: pd.Index
+    user_index_: pd.Index[Any]
     "The index of user IDs."
+    item_index_: pd.Index[Any]
+    "The index of item IDs."
     user_means_: torch.Tensor | None
     "Mean rating for each known user."
     user_vectors_: torch.Tensor
@@ -121,7 +121,10 @@ class UserUser(Predictor):
         self.aggregate = defaults["aggregate"]
         self.use_ratings = defaults["use_ratings"]
 
-    def fit(self, ratings, **kwargs):
+        if self.aggregate not in [self.AGG_WA, self.AGG_SUM]:
+            raise ValueError(f"invalid aggregate {self.aggregate}")
+
+    def fit(self, ratings: pd.DataFrame, **kwargs) -> Self:
         """
         "Train" a user-user CF model.  This memorizes the rating data in a format that is usable
         for future computations.
@@ -146,7 +149,7 @@ class UserUser(Predictor):
         normed, _norms = normalize_sparse_rows(rmat, "unit")
 
         self.user_vectors_ = normed
-        self.user_ratings_ = rmat.to_sparse_coo()
+        self.user_ratings_ = rmat.to_sparse_coo().coalesce()
         self.user_index_ = users
         self.user_means_ = means
         self.item_index_ = items
@@ -169,24 +172,25 @@ class UserUser(Predictor):
         """
 
         watch = util.Stopwatch()
-        items = pd.Index(items, name="item")
+        items = np.asarray(items)
 
         udata = self._get_user_data(user, ratings)
         if udata is None:
             _log.debug("user %s has no ratings, skipping", user)
             return pd.Series(index=items, dtype="float32")
 
-        index, ratings, umean = udata
-        assert len(ratings) == len(self.item_index_)  # ratings is a dense vector
+        uidx, ratings, umean = udata
+        _log.debug("scoring %d items for user %s (idx %d)", len(items), user, uidx or -1)
+        assert ratings.shape == (len(self.item_index_),)  # ratings is a dense vector
 
-        # now ratings is normalized to be a mean-centered unit vector
+        # now ratings has vbeen normalized to be a mean-centered unit vector
         # this means we can dot product to score neighbors
         # score the neighbors!
-        nbr_sims = torch.mv(self.user_vectors_, ratings)
+        nbr_sims = safe_spmv(self.user_vectors_, ratings)
         assert nbr_sims.shape == (len(self.user_index_),)
-        if index is not None:
+        if uidx is not None:
             # zero out the self-similarity
-            nbr_sims[index] = 0
+            nbr_sims[uidx] = 0
 
         # get indices for these neighbors
         nbr_idxs = torch.arange(len(self.user_index_), dtype=torch.int64)
@@ -195,31 +199,35 @@ class UserUser(Predictor):
 
         kn_sims = nbr_sims[nbr_mask]
         kn_idxs = nbr_idxs[nbr_mask]
+        _log.debug(
+            "user %s: %d candidate neighbors (of %d total), max sim %0.4f",
+            user,
+            len(kn_sims),
+            len(self.user_index_),
+            torch.max(kn_sims).item(),
+        )
+        assert not torch.any(torch.isnan(kn_sims))
 
-        _log.debug("found %d candidate neighbor similarities", kn_sims.shape[0])
-
-        iidxs = self.item_index_.get_indexer(items.values)
+        iidxs = self.item_index_.get_indexer_for(items)
         iidxs = torch.from_numpy(iidxs)
-        if self.aggregate == self.AGG_WA:
-            agg = _agg_weighted_avg
-        elif self.aggregate == self.AGG_SUM:
-            agg = _agg_sum
-        else:
-            raise ValueError("invalid aggregate " + self.aggregate)
+
+        ki_mask = iidxs >= 0
+        usable_iidxs = iidxs[ki_mask]
 
         scores = score_items_with_neighbors(
-            iidxs,
+            usable_iidxs,
             kn_idxs,
             kn_sims,
             self.user_ratings_,
             self.nnbrs,
             self.min_nbrs,
-            agg,
+            self.aggregate == "weighted-average",
         )
 
         scores += umean
 
-        results = pd.Series(scores, index=items, name="prediction")
+        results = pd.Series(scores.numpy(), index=items[ki_mask.numpy()], name="prediction")
+        results = results.reindex(items)
 
         _log.debug(
             "scored %d of %d items for %s in %s", results.notna().sum(), len(items), user, watch
@@ -240,8 +248,13 @@ class UserUser(Predictor):
                 _log.warning("user %d has no ratings and none provided", user)
                 return None
 
+            assert index >= 0
             row = self.user_vectors_[index].to_dense()
-            umean = self.user_means_[index].item() if self.user_means_ is not None else 0
+            if self.center:
+                assert self.user_means_ is not None
+                umean = self.user_means_[index].item()
+            else:
+                umean = 0
             return UserRatings(index, row, umean)
         else:
             _log.debug("using provided ratings for user %d", user)
@@ -270,43 +283,6 @@ class UserRatings(NamedTuple):
     mean: float
 
 
-@njit
-def _agg_weighted_avg(iur, item, sims, use):
-    """
-    Weighted-average aggregate.
-
-    Args:
-        iur(matrix._CSR): the item-user ratings matrix
-        item(int): the item index in ``iur``
-        sims(numpy.ndarray): the similarities for the users who have rated ``item``
-        use(numpy.ndarray): positions in sims and the rating row to actually use
-    """
-    rates = iur.row_vs(item)
-    num = 0.0
-    den = 0.0
-    for j in use:
-        num += rates[j] * sims[j]
-        den += np.abs(sims[j])
-    return num / den
-
-
-@njit
-def _agg_sum(iur, item, sims, use):
-    """
-    Sum aggregate
-
-    Args:
-        iur(matrix._CSR): the item-user ratings matrix
-        item(int): the item index in ``iur``
-        sims(numpy.ndarray): the similarities for the users who have rated ``item``
-        use(numpy.ndarray): positions in sims and the rating row to actually use
-    """
-    x = 0.0
-    for j in use:
-        x += sims[j]
-    return x
-
-
 def score_items_with_neighbors(
     items: torch.Tensor,
     nbr_rows: torch.Tensor,
@@ -314,20 +290,29 @@ def score_items_with_neighbors(
     ratings: torch.Tensor,
     max_nbrs: int,
     min_nbrs: int,
-    agg,
+    average: bool,
 ) -> torch.Tensor:
     # select a sub-matrix for further manipulation
     (ni,) = items.shape
     nbr_rates = ratings.index_select(0, nbr_rows)
     nbr_rates = nbr_rates.index_select(1, items)
     nbr_rates = nbr_rates.coalesce()
+    assert nbr_rates.shape == (
+        len(nbr_rows),
+        ni,
+    ), f"nbr rates has shape {nbr_rates.shape} (expected {len(nbr_rows)}, {ni})"
     nbr_t = nbr_rates.transpose(0, 1).to_sparse_csr()
 
     # count nbrs for each item
     counts = nbr_t.crow_indices().diff()
     assert counts.shape == items.shape
 
-    _log.debug("scoring %d items, max count %d", ni, torch.max(counts).item())
+    _log.debug(
+        "scoring %d items, max count %d, nbr shape %s",
+        ni,
+        torch.max(counts).item(),
+        nbr_rates.shape,
+    )
 
     # fast-path items with small neighborhoods
     fp_mask = counts <= max_nbrs
@@ -339,13 +324,16 @@ def score_items_with_neighbors(
     ).coalesce()
     results = torch.mv(nbr_fp.transpose(0, 1), nbr_sims)
 
-    nnz = len(nbr_fp.values())
-    nbr_fp_ones = torch.sparse_coo_tensor(
-        indices=nbr_rates.indices()[:, r_mask],
-        values=torch.ones(nnz),
-        size=nbr_rates.shape,
-    )
-    results /= torch.mv(nbr_fp_ones.transpose(0, 1), nbr_sims)
+    if average:
+        nnz = len(nbr_fp.values())
+        nbr_fp_ones = torch.sparse_coo_tensor(
+            indices=nbr_fp.indices(),
+            values=torch.ones(nnz),
+            size=nbr_rates.shape,
+        )
+        tot_sims = torch.mv(nbr_fp_ones.transpose(0, 1), nbr_sims)
+        assert torch.all(torch.isfinite(tot_sims))
+        results /= tot_sims
 
     # clear out too-small neighborhoods
     results[counts < min_nbrs] = torch.nan
@@ -354,7 +342,9 @@ def score_items_with_neighbors(
     exc_mask = counts > max_nbrs
     if torch.any(exc_mask):
         _log.debug("scoring %d slow-path items", torch.sum(exc_mask))
-    for badi in items[exc_mask]:
+
+    bads = torch.argwhere(exc_mask)[:, 0]
+    for badi in bads:
         col = nbr_t[badi]
         assert col.shape == nbr_rates.shape[:1]
 
@@ -363,42 +353,10 @@ def score_items_with_neighbors(
         bi_sims = nbr_sims[bi_users]
 
         tk_vs, tk_is = torch.topk(bi_sims, max_nbrs)
-        results[badi] = torch.sum(tk_vs * bi_rates[tk_is])
+        sum = torch.sum(tk_vs)
+        if average:
+            results[badi] = torch.dot(tk_vs, bi_rates[tk_is]) / sum
+        else:
+            results[badi] = sum
 
     return results
-
-
-@njit
-def _score(items, results, iur, sims, nnbrs, min_sim, min_nbrs, agg):
-    h_ks = np.empty(nnbrs, dtype=np.int32)
-    h_vs = np.empty(nnbrs)
-    used = np.zeros(len(results), dtype=np.int32)
-
-    for i in range(len(results)):
-        item = items[i]
-        if item < 0:
-            continue
-
-        h_ep = 0
-
-        # who has rated this item?
-        i_users = iur.row_cs(item)
-
-        # what are their similarities to our target user?
-        i_sims = sims[i_users]
-
-        # which of these neighbors do we really want to use?
-        for j, s in enumerate(i_sims):
-            if np.abs(s) < 1.0e-10:
-                continue
-            if min_sim is not None and s < min_sim:
-                continue
-            h_ep = kvp_minheap_insert(0, h_ep, nnbrs, j, s, h_ks, h_vs)
-
-        if h_ep < min_nbrs:
-            continue
-
-        results[i] = agg(iur, item, i_sims, h_ks[:h_ep])
-        used[i] = h_ep
-
-    return used
diff --git a/lenskit/batch/_predict.py b/lenskit/batch/_predict.py
index fc76dd2e..fc53d0ac 100644
--- a/lenskit/batch/_predict.py
+++ b/lenskit/batch/_predict.py
@@ -21,7 +21,7 @@ def _predict_user(model, req):
     res = model.predict_for_user(user, udf["item"])
     res = pd.DataFrame({"user": user, "item": res.index, "prediction": res.values})
     _logger.debug(
-        "%s produced %f/%d predictions for %s in %s",
+        "%s produced %d/%d predictions for %s in %s",
         model,
         res.prediction.notna().sum(),
         len(udf),
diff --git a/pyproject.toml b/pyproject.toml
index 1df3d3cf..8710fcd9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,7 +26,7 @@ dependencies = [
   "numpy >= 1.23",
   "scipy >= 1.9.0",
   "numba >= 0.56, < 0.59",
-  "torch ~=2.1",            # p2c: -s pytorch==2.*
+  "torch ~=2.1",            # p2c: -s pytorch>=2.1,<3
   "threadpoolctl >=3.0",
   "binpickle >= 0.3.2",
   "seedbank >= 0.2.0a2",    # p2c: -p
@@ -56,6 +56,7 @@ test = [
   "pytest-doctestplus ==1.*",
   "pytest-cov >= 2.12",
   "pytest-benchmark ==4.*",
+  "pytest-repeat >=0.9",
   "coverage >= 5",
   "hypothesis >= 6",
 ]
@@ -102,7 +103,7 @@ version_scheme = "release-branch-semver"
 
 # settings for generating conda environments for dev & CI, when needed
 [tool.pyproject2conda]
-channels = ["conda-forge", "pytorch", "nodefaults"]
+channels = ["pytorch", "conda-forge", "nodefaults"]
 python = ["3.10", "3.11"]
 default_envs = ["test", "doc"]
 template_python = "envs/lenskit-py{py_version}-{env}"


Test patch:
diff --git a/envs/lenskit-py3.10-test.yaml b/envs/lenskit-py3.10-test.yaml
index 961ec9ff..d12ddb57 100644
--- a/envs/lenskit-py3.10-test.yaml
+++ b/envs/lenskit-py3.10-test.yaml
@@ -8,8 +8,8 @@
 # Instead edit the corresponding pyproject.toml file.
 #
 channels:
-  - conda-forge
   - pytorch
+  - conda-forge
   - nodefaults
 dependencies:
   - python=3.10
@@ -24,8 +24,9 @@ dependencies:
   - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus==1.*
+  - pytest-repeat>=0.9
   - pytest==7.*
-  - pytorch==2.*
+  - pytorch<3,>=2.1
   - scipy>=1.9.0
   - tbb
   - threadpoolctl>=3.0
diff --git a/envs/lenskit-py3.11-test.yaml b/envs/lenskit-py3.11-test.yaml
index 2af0b477..3135dfd9 100644
--- a/envs/lenskit-py3.11-test.yaml
+++ b/envs/lenskit-py3.11-test.yaml
@@ -8,8 +8,8 @@
 # Instead edit the corresponding pyproject.toml file.
 #
 channels:
-  - conda-forge
   - pytorch
+  - conda-forge
   - nodefaults
 dependencies:
   - python=3.11
@@ -24,8 +24,9 @@ dependencies:
   - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus==1.*
+  - pytest-repeat>=0.9
   - pytest==7.*
-  - pytorch==2.*
+  - pytorch<3,>=2.1
   - scipy>=1.9.0
   - tbb
   - threadpoolctl>=3.0
diff --git a/tests/test_knn_user_user.py b/tests/test_knn_user_user.py
index 54e426a2..d7ff28de 100644
--- a/tests/test_knn_user_user.py
+++ b/tests/test_knn_user_user.py
@@ -11,17 +11,14 @@ from pathlib import Path
 import numpy as np
 import pandas as pd
 import torch
-from scipy.sparse import linalg as spla
 
-from pytest import approx, mark
+from pytest import approx, fail, mark
 
 import lenskit.algorithms.knn.user as knn
 import lenskit.util.test as lktu
 from lenskit.algorithms import Recommender
 from lenskit.util import clone
 
-pytestmark = mark.skip("temporarily disabled while fixing")
-
 _log = logging.getLogger(__name__)
 
 ml_ratings = lktu.ml_test.ratings
@@ -246,25 +243,22 @@ def test_uu_known_preds():
     pairs = known_preds.loc[:, ["user", "item"]]
     _log.info("generating %d known predictions", len(pairs))
 
-    preds = batch.predict(algo, pairs)
+    preds = batch.predict(algo, pairs, n_jobs=1)
     merged = pd.merge(known_preds.rename(columns={"prediction": "expected"}), preds)
     assert len(merged) == len(preds)
     merged["error"] = merged.expected - merged.prediction
-    try:
-        assert not any(merged.prediction.isna() & merged.expected.notna())
-    except AssertionError as e:
+    missing = merged.prediction.isna() & merged.expected.notna()
+    if np.any(missing):
         bad = merged[merged.prediction.isna() & merged.expected.notna()]
         _log.error("%d missing predictions:\n%s", len(bad), bad)
-        raise e
+        fail(f"missing predictions for {np.sum(missing)} items")
 
     err = merged.error
     err = err[err.notna()]
-    try:
-        assert all(err.abs() < 0.01)
-    except AssertionError as e:
-        bad = merged[merged.error.notna() & (merged.error.abs() >= 0.01)]
+    if np.any(err.abs() > 0.01):
+        bad = merged[merged.error.notna() & (merged.error.abs() > 0.01)]
         _log.error("%d erroneous predictions:\n%s", len(bad), bad)
-        raise e
+        fail(f"{len(bad)} erroneous predictions")
 
 
 def __batch_eval(job):
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index 82ff6736..bf6afda9 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -155,4 +155,4 @@ def test_safe_spmv(data):
     tv = torch.from_numpy(v)
 
     tres = safe_spmv(TM, tv)
-    assert tres.cpu().numpy() == approx(res, rel=1.0e-4, abs=1.0e-6)
+    assert tres.cpu().numpy() == approx(res, rel=1.0e-4, abs=1.0e-5)