Repository: pyg-team/pytorch_geometric
Pull request #9246 (2024-04-26 13:10:15)

Fix `test_profiler`


Patch:
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index b2478d75f..8eb174913 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -1256,6 +1256,17 @@ class EdgeIndex(Tensor):
         self._cat_metadata = None
         return self
 
+    def to_vector(self) -> Tensor:
+        r"""Converts :class:`EdgeIndex` into a one-dimensional index
+        vector representation.
+        """
+        num_rows, num_cols = self.get_sparse_size()
+
+        if num_rows * num_cols > torch_geometric.typing.MAX_INT64:
+            raise ValueError("'to_vector()' will result in an overflow")
+
+        return self._data[0] * num_rows + self._data[1]
+
 
 class SortReturnType(NamedTuple):
     values: EdgeIndex
diff --git a/torch_geometric/profile/profiler.py b/torch_geometric/profile/profiler.py
index 9140ecab0..d7eb590d8 100644
--- a/torch_geometric/profile/profiler.py
+++ b/torch_geometric/profile/profiler.py
@@ -331,21 +331,23 @@ def _build_measure_tuple(events: List, occurrences: List) -> NamedTuple:
         hasattr(e, "self_cpu_memory_usage") for e in events)
     if has_self_cpu_memory:
         self_cpu_memory = sum(
-            [getattr(e, "self_cpu_memory_usage", 0) for e in events])
+            [getattr(e, "self_cpu_memory_usage", 0) or 0 for e in events])
     cpu_memory = None
     has_cpu_memory = any(hasattr(e, "cpu_memory_usage") for e in events)
     if has_cpu_memory:
-        cpu_memory = sum([getattr(e, "cpu_memory_usage", 0) for e in events])
+        cpu_memory = sum(
+            [getattr(e, "cpu_memory_usage", 0) or 0 for e in events])
     self_cuda_memory = None
     has_self_cuda_memory = any(
         hasattr(e, "self_cuda_memory_usage") for e in events)
     if has_self_cuda_memory:
         self_cuda_memory = sum(
-            [getattr(e, "self_cuda_memory_usage", 0) for e in events])
+            [getattr(e, "self_cuda_memory_usage", 0) or 0 for e in events])
     cuda_memory = None
     has_cuda_memory = any(hasattr(e, "cuda_memory_usage") for e in events)
     if has_cuda_memory:
-        cuda_memory = sum([getattr(e, "cuda_memory_usage", 0) for e in events])
+        cuda_memory = sum(
+            [getattr(e, "cuda_memory_usage", 0) or 0 for e in events])
 
     # self CUDA time supported in torch >= 1.7
     self_cuda_total = None
@@ -353,13 +355,13 @@ def _build_measure_tuple(events: List, occurrences: List) -> NamedTuple:
         hasattr(e, "self_cuda_time_total") for e in events)
     if has_self_cuda_time:
         self_cuda_total = sum(
-            [getattr(e, "self_cuda_time_total", 0) for e in events])
+            [getattr(e, "self_cuda_time_total", 0) or 0 for e in events])
 
     return Measure(
-        self_cpu_total=sum([e.self_cpu_time_total for e in events]),
-        cpu_total=sum([e.cpu_time_total for e in events]),
+        self_cpu_total=sum([e.self_cpu_time_total or 0 for e in events]),
+        cpu_total=sum([e.cpu_time_total or 0 for e in events]),
         self_cuda_total=self_cuda_total,
-        cuda_total=sum([e.cuda_time_total for e in events]),
+        cuda_total=sum([e.cuda_time_total or 0 for e in events]),
         self_cpu_memory=self_cpu_memory,
         cpu_memory=cpu_memory,
         self_cuda_memory=self_cuda_memory,


Test patch:
diff --git a/test/profile/test_profiler.py b/test/profile/test_profiler.py
index bed81d13f..b476138a6 100644
--- a/test/profile/test_profiler.py
+++ b/test/profile/test_profiler.py
@@ -1,4 +1,3 @@
-import pytest
 import torch
 
 from torch_geometric.nn import GraphSAGE
@@ -7,7 +6,6 @@ from torch_geometric.testing import withDevice
 
 
 @withDevice
-@pytest.mark.skip(reason="Test error")  # TODO Investigate test errors
 def test_profiler(capfd, get_dataset, device):
     x = torch.randn(10, 16, device=device)
     edge_index = torch.tensor([
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index ca045b988..46dbe89ee 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -1088,6 +1088,15 @@ def test_global_mapping(device, dtype):
     assert out.equal(expected)
 
 
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_to_vector(device, dtype):
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], device=device, dtype=dtype)
+    out = adj.to_vector()
+    assert not isinstance(out, EdgeIndex)
+    assert out.equal(tensor([1, 3, 5, 7], device=device))
+
+
 @withCUDA
 @pytest.mark.parametrize('dtype', DTYPES)
 def test_save_and_load(dtype, device, tmp_path):