Repository: pyg-team/pytorch_geometric
Pull request #8533 (2023-12-05 09:12:12)

Introduce `withoutExtensions` decorator


Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 1288e18e7..78046581b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -75,7 +75,7 @@ full = [
     "numba",
     "sympy",
     "pandas",
-    "captum",
+    "captum<0.7.0",
     "rdflib",
     "trimesh",
     "networkx",
diff --git a/torch_geometric/data/edge_index.py b/torch_geometric/data/edge_index.py
index d8024692d..64ad3ec46 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -1251,27 +1251,38 @@ def matmul(
         raise NotImplementedError("`transpose=True` not yet supported for "
                                   "sparse-sparse matrix multiplication")
 
-    if input.is_sorted_by_col:
+    if torch_geometric.typing.WITH_WINDOWS:
+        input = input.to_sparse_coo(input_value)
+    elif input.is_sorted_by_col:
         input = input.to_sparse_csc(input_value)
     else:
         input = input.to_sparse_csr(input_value)
 
-    if other.is_sorted_by_col:
+    if torch_geometric.typing.WITH_WINDOWS:
+        other = other.to_sparse_coo(input_value)
+    elif other.is_sorted_by_col:
         other = other.to_sparse_csc(other_value)
     else:
         other = other.to_sparse_csr(other_value)
 
     out = torch.matmul(input, other)
-    assert out.layout == torch.sparse_csr
 
-    rowptr, col = out.crow_indices(), out.col_indices()
-    edge_index = torch._convert_indices_from_csr_to_coo(
-        rowptr, col, out_int32=rowptr.dtype != torch.int64)
+    rowptr: Optional[Tensor] = None
+    if out.layout == torch.sparse_csr:
+        rowptr, col = out.crow_indices(), out.col_indices()
+        edge_index = torch._convert_indices_from_csr_to_coo(
+            rowptr, col, out_int32=rowptr.dtype != torch.int64)
+
+    elif out.layout == torch.sparse_coo:
+        edge_index = out.indices()
+
+    else:
+        raise NotImplementedError
 
     edge_index = edge_index.as_subclass(EdgeIndex)
     edge_index._sort_order = SortOrder.ROW
     edge_index._sparse_size = (out.size(0), out.size(1))
-    edge_index._rowptr = rowptr
+    edge_index._indptr = rowptr
 
     return edge_index, out.values()
 
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index a10a38a11..67603bfe1 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -1,6 +1,5 @@
 import inspect
 import os
-import platform
 import sys
 import warnings
 from typing import Dict, List, Optional, Tuple, Union
@@ -17,7 +16,6 @@ WITH_PT112 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 12
 WITH_PT113 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 13
 
 WITH_WINDOWS = os.name == 'nt'
-WITH_ARM = platform.machine() != 'x86_64'
 
 if not hasattr(torch, 'sparse_csc'):
     torch.sparse_csc = torch.sparse_coo
diff --git a/torch_geometric/utils/sparse.py b/torch_geometric/utils/sparse.py
index 5f248e1d4..253a511a5 100644
--- a/torch_geometric/utils/sparse.py
+++ b/torch_geometric/utils/sparse.py
@@ -183,8 +183,7 @@ def to_torch_coo_tensor(
         edge_index, edge_attr = coalesce(edge_index, edge_attr, max(size))
 
     if edge_attr is None:
-        if (torch_geometric.typing.WITH_PT20
-                and not torch_geometric.typing.WITH_ARM):
+        if torch_geometric.typing.WITH_PT20 and not edge_index.is_cuda:
             edge_attr = torch.ones(1, device=edge_index.device)
             edge_attr = edge_attr.expand(edge_index.size(1))
         else:
@@ -248,8 +247,7 @@ def to_torch_csr_tensor(
         edge_index, edge_attr = coalesce(edge_index, edge_attr, max(size))
 
     if edge_attr is None:
-        if (torch_geometric.typing.WITH_PT20
-                and not torch_geometric.typing.WITH_ARM):
+        if torch_geometric.typing.WITH_PT20 and not edge_index.is_cuda:
             edge_attr = torch.ones(1, device=edge_index.device)
             edge_attr = edge_attr.expand(edge_index.size(1))
         else:
@@ -318,8 +316,7 @@ def to_torch_csc_tensor(
                                          sort_by_row=False)
 
     if edge_attr is None:
-        if (torch_geometric.typing.WITH_PT20
-                and not torch_geometric.typing.WITH_ARM):
+        if torch_geometric.typing.WITH_PT20 and not edge_index.is_cuda:
             edge_attr = torch.ones(1, device=edge_index.device)
             edge_attr = edge_attr.expand(edge_index.size(1))
         else:


Test patch:
diff --git a/test/conftest.py b/test/conftest.py
index 1a57af8f8..4c61bf43d 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -61,17 +61,32 @@ def get_dataset() -> Callable:
         fs.rm(root)
 
 
+@pytest.fixture
+def enable_extensions():  # Nothing to do.
+    yield
+
+
 @pytest.fixture
 def disable_extensions():
-    prev_state = {
-        'WITH_PYG_LIB': torch_geometric.typing.WITH_PYG_LIB,
-        'WITH_SAMPLED_OP': torch_geometric.typing.WITH_SAMPLED_OP,
-        'WITH_INDEX_SORT': torch_geometric.typing.WITH_INDEX_SORT,
-        'WITH_TORCH_SCATTER': torch_geometric.typing.WITH_TORCH_SCATTER,
-        'WITH_TORCH_SPARSE': torch_geometric.typing.WITH_TORCH_SPARSE,
-    }
-    for key in prev_state.keys():
+    def is_setting(name: str) -> bool:
+        if not name.startswith('WITH_'):
+            return False
+        if name.startswith('WITH_PT') or name.startswith('WITH_WINDOWS'):
+            return False
+        return True
+
+    settings = dir(torch_geometric.typing)
+    settings = [key for key in settings if is_setting(key)]
+    state = {key: getattr(torch_geometric.typing, key) for key in settings}
+
+    for key in state.keys():
         setattr(torch_geometric.typing, key, False)
     yield
-    for key, value in prev_state.items():
+    for key, value in state.items():
         setattr(torch_geometric.typing, key, value)
+
+
+@pytest.fixture
+def without_extensions(request):
+    request.getfixturevalue(request.param)
+    return request.param == 'disable_extensions'
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 2388213bd..d72c9997b 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -1,6 +1,6 @@
 import os.path as osp
 import warnings
-from typing import Optional
+from typing import List, Optional
 
 import pytest
 import torch
@@ -22,6 +22,7 @@ from torch_geometric.testing import (
     onlyCUDA,
     onlyLinux,
     withCUDA,
+    withoutExtensions,
     withPackage,
 )
 from torch_geometric.typing import SparseTensor
@@ -702,10 +703,14 @@ def test_torch_spmm(device, reduce, transpose, is_undirected):
 
 
 @withCUDA
+@withoutExtensions
 @pytest.mark.parametrize('reduce', ReduceType.__args__)
 @pytest.mark.parametrize('transpose', TRANSPOSE)
 @pytest.mark.parametrize('is_undirected', IS_UNDIRECTED)
-def test_spmm(device, reduce, transpose, is_undirected):
+def test_spmm(without_extensions, device, reduce, transpose, is_undirected):
+    if without_extensions:
+        warnings.filterwarnings('ignore', '.*can be accelerated via.*')
+
     if is_undirected:
         kwargs = dict(is_undirected=True)
         adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], device=device, **kwargs)
@@ -768,7 +773,8 @@ def test_spspmm():
     assert isinstance(out, EdgeIndex)
     assert out.is_sorted_by_row
     assert out._sparse_size == (3, 3)
-    assert out._rowptr is not None
+    if not torch_geometric.typing.WITH_WINDOWS:
+        assert out._indptr is not None
     assert torch.allclose(out.to_dense(value), adj1_dense @ adj1_dense)
 
     out, value = adj1 @ adj2
@@ -785,7 +791,8 @@ def test_spspmm():
 
 
 @withCUDA
-def test_matmul(device):
+@withoutExtensions
+def test_matmul(without_extensions, device):
     kwargs = dict(sort_order='row', device=device)
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], **kwargs)
     x = torch.randn(3, 1, device=device)
@@ -832,6 +839,10 @@ def test_save_and_load(dtype, device, tmp_path):
     assert out._indptr.equal(adj._indptr)
 
 
+def _collate_fn(edge_indices: List[EdgeIndex]) -> List[EdgeIndex]:
+    return edge_indices
+
+
 @pytest.mark.parametrize('dtype', DTYPES)
 @pytest.mark.parametrize('num_workers', [0, 2])
 def test_data_loader(dtype, num_workers):
@@ -843,7 +854,7 @@ def test_data_loader(dtype, num_workers):
         [adj] * 4,
         batch_size=2,
         num_workers=num_workers,
-        collate_fn=lambda x: x,
+        collate_fn=_collate_fn,
         drop_last=True,
     )
 
diff --git a/test/io/test_fs.py b/test/io/test_fs.py
index 0e92d112e..48b56cf34 100644
--- a/test/io/test_fs.py
+++ b/test/io/test_fs.py
@@ -6,6 +6,7 @@ import pytest
 
 from torch_geometric.data import extract_zip
 from torch_geometric.io import fs
+from torch_geometric.testing import noWindows
 
 
 @pytest.fixture(params=['file', 'memory'])
@@ -22,6 +23,7 @@ def test_get_fs():
     assert 'memory' in fs.get_fs('memory:///tmp/test').protocol
 
 
+@noWindows
 def test_normpath():
     assert fs.normpath('////home') == '/home'
     assert fs.normpath('memory:////home') == 'memory:////home'
diff --git a/torch_geometric/testing/__init__.py b/torch_geometric/testing/__init__.py
index 58eaded74..509362fa0 100644
--- a/torch_geometric/testing/__init__.py
+++ b/torch_geometric/testing/__init__.py
@@ -18,6 +18,7 @@ from .decorators import (
     withPackage,
     withCUDA,
     disableExtensions,
+    withoutExtensions,
 )
 from .asserts import assert_module
 from .feature_store import MyFeatureStore
@@ -43,6 +44,7 @@ __all__ = [
     'withPackage',
     'withCUDA',
     'disableExtensions',
+    'withoutExtensions',
     'assert_module',
     'MyFeatureStore',
     'MyGraphStore',
diff --git a/torch_geometric/testing/decorators.py b/torch_geometric/testing/decorators.py
index b2da3a5c1..f32d775b5 100644
--- a/torch_geometric/testing/decorators.py
+++ b/torch_geometric/testing/decorators.py
@@ -199,3 +199,17 @@ def disableExtensions(func: Callable):
     import pytest
 
     return pytest.mark.usefixtures('disable_extensions')(func)
+
+
+def withoutExtensions(func: Callable):
+    r"""A decorator to test both with and without the usage of extension
+    packages such as :obj:`torch_scatter`, :obj:`torch_sparse` and
+    :obj:`pyg_lib`.
+    """
+    import pytest
+
+    return pytest.mark.parametrize(
+        'without_extensions',
+        ['enable_extensions', 'disable_extensions'],
+        indirect=True,
+    )(func)