missing op `grid_cluster` on CPU

[christian-rauch]

🐛 Describe the bug

I installed the CPU version of pyg-lib via:

pip install pyg-lib -f https://data.pyg.org/whl/torch-2.13.0%2Bcpu.html

and when running torch.ops.pyg.grid_cluster, I get:

Could not run 'pyg::grid_cluster' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'pyg::grid_cluster' is only available for these backends: [CPU, Meta, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradMPS, AutogradXPU, AutogradHPU, AutogradLazy, AutogradMTIA, AutogradMAIA, AutogradPrivateUse1, AutogradMeta, Tracer, AutocastCPU, AutocastMTIA, AutocastMAIA, AutocastXPU, AutocastMPS, AutocastCUDA, FuncTorchBatched, BatchedNestedTensor, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PreDispatch, PythonDispatcher].

CPU: registered at /project/pyg_lib/csrc/ops/cpu/cluster_kernel.cpp:48 [kernel]
Meta: registered at /__w/pytorch/pytorch/aten/src/ATen/core/MetaFallbackKernel.cpp:23 [backend fallback]
BackendSelect: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
Python: registered at /__w/pytorch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:192 [backend fallback]
FuncTorchDynamicLayerBackMode: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:479 [backend fallback]
Functionalize: registered at /__w/pytorch/pytorch/aten/src/ATen/FunctionalizeFallbackKernel.cpp:411 [backend fallback]
Conjugate: registered at /__w/pytorch/pytorch/aten/src/ATen/ConjugateFallback.cpp:17 [backend fallback]
Negative: registered at /__w/pytorch/pytorch/aten/src/ATen/native/NegateFallback.cpp:18 [backend fallback]
ZeroTensor: registered at /__w/pytorch/pytorch/aten/src/ATen/ZeroTensorFallback.cpp:115 [backend fallback]
ADInplaceOrView: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:103 [backend fallback]
AutogradOther: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:62 [backend fallback]
AutogradCPU: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:66 [backend fallback]
AutogradCUDA: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:74 [backend fallback]
AutogradXLA: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:86 [backend fallback]
AutogradMPS: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:94 [backend fallback]
AutogradXPU: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:70 [backend fallback]
AutogradHPU: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:107 [backend fallback]
AutogradLazy: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:90 [backend fallback]
AutogradMTIA: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:78 [backend fallback]
AutogradMAIA: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:82 [backend fallback]
AutogradPrivateUse1: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:111 [backend fallback]
AutogradMeta: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:98 [backend fallback]
Tracer: registered at /__w/pytorch/pytorch/torch/csrc/autograd/TraceTypeManual.cpp:296 [backend fallback]
AutocastCPU: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:330 [backend fallback]
AutocastMTIA: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:472 [backend fallback]
AutocastMAIA: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:510 [backend fallback]
AutocastXPU: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:548 [backend fallback]
AutocastMPS: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:221 [backend fallback]
AutocastCUDA: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:177 [backend fallback]
FuncTorchBatched: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp:730 [backend fallback]
BatchedNestedTensor: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp:757 [backend fallback]
FuncTorchVmapMode: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/VmapModeRegistrations.cpp:22 [backend fallback]
Batched: registered at /__w/pytorch/pytorch/aten/src/ATen/LegacyBatchingRegistrations.cpp:1072 [backend fallback]
VmapMode: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/VmapModeRegistrations.cpp:32 [backend fallback]
FuncTorchGradWrapper: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/TensorWrapper.cpp:210 [backend fallback]
PythonTLSSnapshot: registered at /__w/pytorch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:200 [backend fallback]
FuncTorchDynamicLayerFrontMode: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:475 [backend fallback]
PreDispatch: registered at /__w/pytorch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:204 [backend fallback]
PythonDispatcher: registered at /__w/pytorch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:196 [backend fallback]

Environment

• pyg-lib version: 0.8.0+pt213cpu
• PyTorch version: 2.13.0+rocm7.2
• OS: Linux 7.0.0-27-generic
• Python version: Python 3.14.4
• CUDA/cuDNN version: none
• How you installed PyTorch and pyg-lib (conda, pip, source): pip
• Any other relevant information:

[akihironitta]

Hi @christian-rauch, Thank you for reporting this issue. It looks like you're running it with a ROCm build of PyTorch 2.13 and a CPU build of pyg-lib, and the error messaging suggests you're passing CUDA (ROCm) tensor to the op. Could you share your Python script to reproduce this error?

[christian-rauch]

Indeed. I am using this on ROCm. But since ROCm is not supported in pyg-lib (#711), I am using the CPU version. I would expect that tensors are moved to the CPU device, when they are used by pyg-lib.

This example usage of voxel_grid demonstrates the problem:

import torch
from torch_geometric.nn.pool import voxel_grid

print(f"PyTorch:      {torch.__version__}")
print(f"CUDA version: {torch.version.cuda}")
print(f"ROCm (HIP):   {torch.version.hip}")

pos = torch.tensor([
    [0.1, 0.2, 0.1],
    [0.2, 0.1, 0.3],
    [1.1, 0.9, 1.0],
    [1.2, 1.1, 0.8],
    [2.5, 2.5, 2.5],
], dtype=torch.float, device=torch.device("cuda"))

cluster = voxel_grid(pos, size=1.0)

print(cluster)

Output:

PyTorch:      2.13.0+rocm7.2
CUDA version: None
ROCm (HIP):   7.2.53211
Traceback (most recent call last):
  File "/home/christian/dev/voxel_grid_example.py", line 16, in <module>
    cluster = voxel_grid(pos, size=1.0)
  File "/home/christian/.local/lib/python3.14/site-packages/torch_geometric/nn/pool/voxel_grid.py", line 66, in voxel_grid
    return torch.ops.pyg.grid_cluster(pos, size, start, end)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/christian/.local/lib/python3.14/site-packages/torch/_ops.py", line 1279, in __call__
    return self._op(*args, **kwargs)
           ~~~~~~~~^^^^^^^^^^^^^^^^^
NotImplementedError: Could not run 'pyg::grid_cluster' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'pyg::grid_cluster' is only available for these backends: [CPU, Meta, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradMPS, AutogradXPU, AutogradHPU, AutogradLazy, AutogradMTIA, AutogradMAIA, AutogradPrivateUse1, AutogradMeta, Tracer, AutocastCPU, AutocastMTIA, AutocastMAIA, AutocastXPU, AutocastMPS, AutocastCUDA, FuncTorchBatched, BatchedNestedTensor, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PreDispatch, PythonDispatcher].

CPU: registered at /project/pyg_lib/csrc/ops/cpu/cluster_kernel.cpp:48 [kernel]
Meta: registered at /__w/pytorch/pytorch/aten/src/ATen/core/MetaFallbackKernel.cpp:23 [backend fallback]
BackendSelect: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
Python: registered at /__w/pytorch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:192 [backend fallback]
FuncTorchDynamicLayerBackMode: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:479 [backend fallback]
Functionalize: registered at /__w/pytorch/pytorch/aten/src/ATen/FunctionalizeFallbackKernel.cpp:411 [backend fallback]
Conjugate: registered at /__w/pytorch/pytorch/aten/src/ATen/ConjugateFallback.cpp:17 [backend fallback]
Negative: registered at /__w/pytorch/pytorch/aten/src/ATen/native/NegateFallback.cpp:18 [backend fallback]
ZeroTensor: registered at /__w/pytorch/pytorch/aten/src/ATen/ZeroTensorFallback.cpp:115 [backend fallback]
ADInplaceOrView: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:103 [backend fallback]
AutogradOther: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:62 [backend fallback]
AutogradCPU: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:66 [backend fallback]
AutogradCUDA: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:74 [backend fallback]
AutogradXLA: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:86 [backend fallback]
AutogradMPS: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:94 [backend fallback]
AutogradXPU: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:70 [backend fallback]
AutogradHPU: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:107 [backend fallback]
AutogradLazy: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:90 [backend fallback]
AutogradMTIA: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:78 [backend fallback]
AutogradMAIA: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:82 [backend fallback]
AutogradPrivateUse1: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:111 [backend fallback]
AutogradMeta: registered at /__w/pytorch/pytorch/aten/src/ATen/core/VariableFallbackKernel.cpp:98 [backend fallback]
Tracer: registered at /__w/pytorch/pytorch/torch/csrc/autograd/TraceTypeManual.cpp:296 [backend fallback]
AutocastCPU: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:330 [backend fallback]
AutocastMTIA: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:472 [backend fallback]
AutocastMAIA: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:510 [backend fallback]
AutocastXPU: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:548 [backend fallback]
AutocastMPS: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:221 [backend fallback]
AutocastCUDA: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/autocast_mode.cpp:177 [backend fallback]
FuncTorchBatched: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp:730 [backend fallback]
BatchedNestedTensor: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/LegacyBatchingRegistrations.cpp:757 [backend fallback]
FuncTorchVmapMode: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/VmapModeRegistrations.cpp:22 [backend fallback]
Batched: registered at /__w/pytorch/pytorch/aten/src/ATen/LegacyBatchingRegistrations.cpp:1072 [backend fallback]
VmapMode: fallthrough registered at /__w/pytorch/pytorch/aten/src/ATen/VmapModeRegistrations.cpp:32 [backend fallback]
FuncTorchGradWrapper: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/TensorWrapper.cpp:210 [backend fallback]
PythonTLSSnapshot: registered at /__w/pytorch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:200 [backend fallback]
FuncTorchDynamicLayerFrontMode: registered at /__w/pytorch/pytorch/aten/src/ATen/functorch/DynamicLayer.cpp:475 [backend fallback]
PreDispatch: registered at /__w/pytorch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:204 [backend fallback]
PythonDispatcher: registered at /__w/pytorch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp:196 [backend fallback]