Fix macOS ARM tests (#9020)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
pyg-team · Mar 5, 2024 · 0d30e89 · 0d30e89
1 parent db9d8f5
commit 0d30e89
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 6 deletions.
diff --git a/test/loader/test_dataloader.py b/test/loader/test_dataloader.py
@@ -11,6 +11,7 @@
 from torch_geometric.testing import (
     get_random_edge_index,
     get_random_tensor_frame,
+    onlyLinux,
     withCUDA,
     withPackage,
 )
@@ -72,6 +73,7 @@ def test_dataloader(num_workers, device):
         assert batch.edge_index_batch.tolist() == [0, 0, 0, 0, 1, 1, 1, 1]
 
 
+@onlyLinux
 @pytest.mark.parametrize('num_workers', num_workers_list)
 def test_dataloader_on_disk_dataset(tmp_path, num_workers):
     dataset = OnDiskDataset(tmp_path)

diff --git a/test/test_edge_index.py b/test/test_edge_index.py
@@ -952,7 +952,7 @@ def test_spspmm(device, reduce, transpose, is_undirected):
         assert isinstance(out, EdgeIndex)
         assert out.is_sorted_by_row
         assert out._sparse_size == (3, 3)
-        if not torch_geometric.typing.WITH_WINDOWS:
+        if torch_geometric.typing.WITH_MKL:
             assert out._indptr is not None
         assert torch.allclose(out.to_dense(value), adj1_dense @ adj2_dense)
     else:

diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
@@ -1823,7 +1823,7 @@ def matmul(
 
     transpose &= not input.is_undirected or input_value is not None
 
-    if torch_geometric.typing.WITH_WINDOWS:  # pragma: no cover
+    if not torch_geometric.typing.WITH_MKL:  # pragma: no cover
         sparse_input = input.to_sparse_coo(input_value)
     elif input.is_sorted_by_col:
         sparse_input = input.to_sparse_csc(input_value)
@@ -1833,7 +1833,7 @@ def matmul(
     if transpose:
         sparse_input = sparse_input.t()
 
-    if torch_geometric.typing.WITH_WINDOWS:  # pragma: no cover
+    if not torch_geometric.typing.WITH_MKL:  # pragma: no cover
         other = other.to_sparse_coo(other_value)
     elif other.is_sorted_by_col:
         other = other.to_sparse_csc(other_value)

diff --git a/torch_geometric/transforms/add_positional_encoding.py b/torch_geometric/transforms/add_positional_encoding.py
@@ -153,10 +153,10 @@ def forward(self, data: Data) -> Data:
             adj = torch.zeros((N, N), device=row.device)
             adj[row, col] = value
             loop_index = torch.arange(N, device=row.device)
-        elif torch_geometric.typing.WITH_WINDOWS:
-            adj = to_torch_coo_tensor(data.edge_index, value, size=data.size())
-        else:
+        elif torch_geometric.typing.WITH_MKL:
             adj = to_torch_csr_tensor(data.edge_index, value, size=data.size())
+        else:
+            adj = to_torch_coo_tensor(data.edge_index, value, size=data.size())
 
         def get_pe(out: Tensor) -> Tensor:
             if is_torch_sparse_tensor(out):

diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
@@ -17,6 +17,7 @@
 WITH_PT113 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 13
 
 WITH_WINDOWS = os.name == 'nt'
+WITH_MKL = 'USE_MKL=OFF' not in torch.__config__.show()
 
 MAX_INT64 = torch.iinfo(torch.int64).max