Repository: pyg-team/pytorch_geometric (python)

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

Pull request:
Introduce dense eigenvalue computation in `AddLaplacianEigenvectorPE` for small-scale graphs
Guarantees that eigenvector computation is stable.

Fixes #8131 

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 114a8dafd..50c2d94ab 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -117,6 +117,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Enabled dense eigenvalue computation in `AddLaplacianEigenvectorPE` for small-scale graphs ([#8143](https://github.com/pyg-team/pytorch_geometric/pull/8143))
 - Fix `DynamicBatchSampler.__len__` to raise an error in case `num_steps` is undefined ([#8137](https://github.com/pyg-team/pytorch_geometric/pull/8137))
 - Enabled pickling of `DimeNet` models ([#8019](https://github.com/pyg-team/pytorch_geometric/pull/8019))
 - Changed the `trim_to_layer` function to filter out non-reachable node and edge types when operating on heterogeneous graphs ([#7942](https://github.com/pyg-team/pytorch_geometric/pull/7942))
diff --git a/torch_geometric/transforms/add_positional_encoding.py b/torch_geometric/transforms/add_positional_encoding.py
index 706cf6621..1524562a3 100644
--- 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 @@ class AddLaplacianEigenvectorPE(BaseTransform):
         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 @@ class AddLaplacianEigenvectorPE(BaseTransform):
         )
 
         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
index a3cd16ce9..c635f9f3d 100644
--- 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.