Merge pull request #12 from denix56/triton

Fix flake8 errors

Author: denix56

torch_cluster/fps.py

@@ -5,32 +5,6 @@
 
 import torch_cluster.typing
 
-@torch.library.register_fake("torch_cluster::fps")
-def _(src, ptr, ratio, random_start = True):
-    torch._check(src.device == ptr.device)
-    torch._check(ptr.ndim == 1)
-
-    ctx = torch.library.get_ctx()
-    nnz = ctx.new_dynamic_size()
-    return ptr.new_empty((nnz,))
-
-@torch.library.register_fake("torch_cluster::fps")
-def _(src, ptr, ratio, random_start=True):
-    torch._check(src.device == ptr.device)
-    torch._check(ptr.ndim == 1)
-
-    ctx = torch.library.get_ctx()
-    nnz = ctx.new_dynamic_size()
-    return ptr.new_empty((nnz,))
-
-@torch.library.register_fake("torch_cluster::fps")
-def _(src, ptr, ratio, random_start=True):
-    torch._check(src.device == ptr.device)
-    torch._check(ptr.ndim == 1)
-
-    ctx = torch.library.get_ctx()
-    nnz = ctx.new_dynamic_size()
-    return ptr.new_empty((nnz,))
 
 @torch.library.register_fake("torch_cluster::fps")
 def _(src, ptr, ratio, random_start=True):

torch_cluster/triton/_kernels.py

@@ -106,9 +106,21 @@ def _nearest_kernel(
     best_idx = tl.zeros((BLOCK_N,), dtype=tl.int32)  # Best y index.
 
     if USE_BATCH:
-        batch_id = tl.load(batch_x_ptr + offs_n, mask=mask_x, other=0)  # Batch id per x.
-        left = tl.load(ptr_y_ptr + batch_id, mask=mask_x, other=0)  # y range start.
-        right = tl.load(ptr_y_ptr + batch_id + 1, mask=mask_x, other=0)  # y range end.
+        batch_id = tl.load(
+            batch_x_ptr + offs_n,
+            mask=mask_x,
+            other=0,
+        )  # Batch id per x.
+        left = tl.load(
+            ptr_y_ptr + batch_id,
+            mask=mask_x,
+            other=0,
+        )  # y range start.
+        right = tl.load(
+            ptr_y_ptr + batch_id + 1,
+            mask=mask_x,
+            other=0,
+        )  # y range end.
     else:
         left = 0  # Full y range.
         right = M
@@ -194,7 +206,11 @@ def _nearest_kernel(
             inv_y = tl.rsqrt(y_sq + EPS)  # 1/||y||.
             dist = 1.0 - acc * (inv_y[:, None] * inv_x[None, :])
         else:
-            dist = tl.fma(-2.0, acc, y_sq[:, None] + x_sq[None, :])  # L2^2 distance.
+            dist = tl.fma(
+                -2.0,
+                acc,
+                y_sq[:, None] + x_sq[None, :],
+            )  # L2^2 distance.
 
         if full_y:
             valid = tl.broadcast_to(mask_x[None, :], (BLOCK_M, BLOCK_N))
@@ -218,7 +234,11 @@ def _nearest_kernel(
         best_dist = tl.where(better, block_min, best_dist)
         best_idx = tl.where(better, block_idx, best_idx)
 
-    tl.store(out_ptr + offs_n, best_idx.to(tl.int64), mask=mask_x)  # Write output.
+    tl.store(
+        out_ptr + offs_n,
+        best_idx.to(tl.int64),
+        mask=mask_x,
+    )  # Write output.
 
 
 @triton.autotune(

torch_cluster/triton/knn.py

@@ -1,6 +1,4 @@
 from __future__ import annotations
-import math
-import triton
 from typing import Optional
 
 import triton