Merge pull request #9 from denix56/triton

Triton

Author: denix56

torch_cluster/fps.py

@@ -32,6 +32,15 @@ def _(src, ptr, ratio, random_start=True):
     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,))
+
 
 def fps(  # noqa
     src: torch.Tensor,