Fix numerical errors

Author: denix56

benchmarks/test_benchmark_knn.py

@@ -52,7 +52,7 @@ def _assert_knn_within_cuda(
     tol=None,
 ):
     if tol is None:
-        tol = 2 * torch.finfo(x.dtype).eps
+        tol = 5 * torch.finfo(x.dtype).eps
     m = y.size(0)
     cuda_rows = out_cuda[0]
     cuda_cols = out_cuda[1]

benchmarks/test_benchmark_nearest.py

@@ -41,7 +41,7 @@ def _assert_nearest_within_cuda(
     tol: float | None = None,
 ) -> None:
     if tol is None:
-        tol = 2 * torch.finfo(x.dtype).eps
+        tol = 5 * torch.finfo(x.dtype).eps
     x_f = x.float()
     y_f = y.float()
     cuda_dist = ((x_f - y_f[out_cuda]) ** 2).sum(dim=1)

benchmarks/test_benchmark_radius.py

@@ -53,7 +53,7 @@ def _assert_radius_within_cuda(
     if edge_index.numel() == 0:
         return
     if tol is None:
-        tol = torch.finfo(x.dtype).eps
+        tol = 5 * torch.finfo(x.dtype).eps
     row, col = edge_index
     x_f = x.float()
     y_f = y.float()

torch_cluster/triton/_kernels.py

@@ -147,6 +147,7 @@ def _nearest_kernel(
         )
         x_sq = x_t
         x_block_ptr_sq = tl.advance(x_block_ptr_sq, (BLOCK_K, 0))
+    x_sq += x_sq_c
     if COSINE:
         inv_x = tl.rsqrt(x_sq + EPS)  # 1/||x||.
 
@@ -220,13 +221,15 @@ def _nearest_kernel(
             y_block_ptr = tl.advance(y_block_ptr, (0, BLOCK_K))
             x_block_ptr = tl.advance(x_block_ptr, (BLOCK_K, 0))
 
+        acc += acc_c
+        y_sq += y_sq_c
         if COSINE:
-            inv_y = tl.rsqrt(y_sq + y_sq_c + EPS)  # 1/||y||.
+            inv_y = tl.rsqrt(y_sq + EPS)  # 1/||y||.
             dist = 1.0 - acc * (inv_y[:, None] * inv_x[None, :])
         else:
             dist = (
-                (y_sq + y_sq_c)[:, None]
-                + (x_sq + x_sq_c)[None, :]
+                (y_sq)[:, None]
+                + (x_sq)[None, :]
                 - 2.0 * acc
             )