Repository: pyg-team/pytorch_geometric

Issue:
`test/transforms/test_add_positional_encoding.py::test_eigenvector_permutation_invariance` fails with AMD CPU
### 🐛 Describe the bug

The unit test added in #8087 fails with the error bellow on AMD CPU based systems.


```
NVIDIA_TF32_OVERRIDE=0 pytest -s --cov --cov-report=xml test/transforms/test_add_positional_encoding.py::test_eigenvector_permutation_invariance
```

Traceback:
```
======================================================================== FAILURES ========================================================================$
_________________________________________________________ test_eigenvector_permutation_invariance ________________________________________________________$

    def test_eigenvector_permutation_invariance():
        edge_index = torch.tensor([[0, 1, 0, 4, 1, 4, 2, 3, 3, 5],
                                   [1, 0, 4, 0, 4, 1, 3, 2, 5, 3]])
        data = Data(edge_index=edge_index, num_nodes=6)

        perm = torch.tensor([5, 4, 3, 2, 1, 0])
        transform = AddLaplacianEigenvectorPE(
            k=1,
            is_undirected=True,
            attr_name='x',
            v0=torch.arange(data.num_nodes),
        )
        out1 = transform(data)

        transform = AddLaplacianEigenvectorPE(
            k=1,
            is_undirected=True,
            attr_name='x',
            v0=perm,
        )
        out2 = transform(data.subgraph(perm))

        print(out1.x[perm].abs())
        print(out2.x.abs())

>       assert torch.allclose(out1.x[perm].abs(), out2.x.abs(), atol=1e-1)
E       assert False
E        +  where False = <built-in method allclose of type object at 0x7f2f94b318a0>(tensor([[0.1360],\n        [0.5556],\n        [0.1924],\n        [0.1
360],\n        [0.5556],\n        [0.5556]]), tensor([[0.2091],\n        [0.5245],\n        [0.2957],\n        [0.2091],\n        [0.5245],\n        [0.524
5]]), atol=0.1)
E        +    where <built-in method allclose of type object at 0x7f2f94b318a0> = torch.allclose
E        +    and   tensor([[0.1360],\n        [0.5556],\n        [0.1924],\n        [0.1360],\n        [0.5556],\n        [0.5556]]) = <built-in method ab
s of Tensor object at 0x7f2bc6e1e160>()
E        +      where <built-in method abs of Tensor object at 0x7f2bc6e1e160> = tensor([[ 0.1360],\n        [-0.5556],\n        [ 0.1924],\n        [ 0.13
60],\n        [-0.5556],\n        [-0.5556]]).abs
E        +    and   tensor([[0.2091],\n        [0.5245],\n        [0.2957],\n        [0.2091],\n        [0.5245],\n        [0.5245]]) = <built-in method ab
s of Tensor object at 0x7f2bc89d3470>()
E        +      where <built-in method abs of Tensor object at 0x7f2bc89d3470> = tensor([[ 0.2091],\n        [-0.5245],\n        [ 0.2957],\n        [ 0.20
91],\n        [-0.5245],\n        [-0.5245]]).abs
E        +        where tensor([[ 0.2091],\n        [-0.5245],\n        [ 0.2957],\n        [ 0.2091],\n        [-0.5245],\n        [-0.5245]]) = Data(edge
_index=[2, 10], num_nodes=6, x=[6, 1]).x
```

### Environment

* PyG version: 2.4.0 (built from source)
* PyTorch version: 2.1.0
* OS:  Ubuntu 22.04.3 LTS
* Python version: 3.10.12
* CUDA/cuDNN version: 12.2
* How you installed PyTorch and PyG (`conda`, `pip`, source): source


Discussion:
Is it just flaky or does it always fail? I don't see the test case use GPU so I wondered why this happens on an A100 particularly.
> Is it just flaky or does it always fail? I don't see the test case use GPU so I wondered why this happens on an A100 particularly.

You are right, it's a CPU only test, it was passing on our V100 machines and failing on A100 ones so I wrongly assumed it has something to do with the GPU.

So it always fails on certain machines but always pass on others.
After more runs across more machines it seems that it's a Intel vs AMD issue. 

PASS on:
```
Architecture:            x86_64
  CPU op-mode(s):        32-bit, 64-bit
  Address sizes:         46 bits physical, 48 bits virtual
  Byte Order:            Little Endian
CPU(s):                  80
  On-line CPU(s) list:   0-79
Vendor ID:               GenuineIntel
  Model name:            Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz
    CPU family:          6
    Model:               79
    Thread(s) per core:  2
    Core(s) per socket:  20
    Socket(s):           2
```

FAIL on:
```
Architecture:            x86_64
  CPU op-mode(s):        32-bit, 64-bit
  Address sizes:         43 bits physical, 48 bits virtual
  Byte Order:            Little Endian
CPU(s):                  256
  On-line CPU(s) list:   0-255
Vendor ID:               AuthenticAMD
  Model name:            AMD EPYC 7742 64-Core Processor
    CPU family:          23
    Model:               49
    Thread(s) per core:  2
    Core(s) per socket:  64
    Socket(s):           2
```
Changing the title and description accordingly.

Resolving patch:
diff --git a/torch_geometric/transforms/add_positional_encoding.py b/torch_geometric/transforms/add_positional_encoding.py
--- a/torch_geometric/transforms/add_positional_encoding.py
+++ b/torch_geometric/transforms/add_positional_encoding.py
@@ -52,6 +52,9 @@ class AddLaplacianEigenvectorPE(BaseTransform):
             :obj:`False`) or :meth:`scipy.sparse.linalg.eigsh` (when
             :attr:`is_undirected` is :obj:`True`).
     """
+    # Number of nodes from which to use sparse eigenvector computation:
+    SPARSE_THRESHOLD: int = 100
+
     def __init__(
         self,
         k: int,
@@ -65,9 +68,6 @@ def __init__(
         self.kwargs = kwargs
 
     def forward(self, data: Data) -> Data:
-        from scipy.sparse.linalg import eigs, eigsh
-        eig_fn = eigs if not self.is_undirected else eigsh
-
         num_nodes = data.num_nodes
         edge_index, edge_weight = get_laplacian(
             data.edge_index,
@@ -77,15 +77,23 @@ def forward(self, data: Data) -> Data:
         )
 
         L = to_scipy_sparse_matrix(edge_index, edge_weight, num_nodes)
-        L = L.tocsr()
-
-        eig_vals, eig_vecs = eig_fn(
-            L,
-            k=self.k + 1,
-            which='SR' if not self.is_undirected else 'SA',
-            return_eigenvectors=True,
-            **self.kwargs,
-        )
+
+        if num_nodes < self.SPARSE_THRESHOLD:
+            from numpy.linalg import eig, eigh
+            eig_fn = eig if not self.is_undirected else eigh
+
+            eig_vals, eig_vecs = eig_fn(L.todense())
+        else:
+            from scipy.sparse.linalg import eigs, eigsh
+            eig_fn = eigs if not self.is_undirected else eigsh
+
+            eig_vals, eig_vecs = eig_fn(
+                L,
+                k=self.k + 1,
+                which='SR' if not self.is_undirected else 'SA',
+                return_eigenvectors=True,
+                **self.kwargs,
+            )
 
         eig_vecs = np.real(eig_vecs[:, eig_vals.argsort()])
         pe = torch.from_numpy(eig_vecs[:, 1:self.k + 1])


Test patch:
diff --git a/test/transforms/test_add_positional_encoding.py b/test/transforms/test_add_positional_encoding.py
--- a/test/transforms/test_add_positional_encoding.py
+++ b/test/transforms/test_add_positional_encoding.py
@@ -56,24 +56,22 @@ def test_eigenvector_permutation_invariance():
                                [1, 0, 4, 0, 4, 1, 3, 2, 5, 3]])
     data = Data(edge_index=edge_index, num_nodes=6)
 
-    perm = torch.tensor([5, 4, 3, 2, 1, 0])
+    perm = torch.randperm(data.num_nodes)
     transform = AddLaplacianEigenvectorPE(
-        k=1,
+        k=2,
         is_undirected=True,
         attr_name='x',
-        v0=torch.arange(data.num_nodes),
     )
     out1 = transform(data)
 
     transform = AddLaplacianEigenvectorPE(
-        k=1,
+        k=2,
         is_undirected=True,
         attr_name='x',
-        v0=perm,
     )
     out2 = transform(data.subgraph(perm))
 
-    assert torch.allclose(out1.x[perm].abs(), out2.x.abs(), atol=1e-1)
+    assert torch.allclose(out1.x[perm].abs(), out2.x.abs(), atol=1e-6)
 
 
 @onlyLinux  # TODO  (matthias) Investigate CSR @ CSR support on Windows.