Clean-up imports

benchmark/utils/__init__.py

@@ -4,7 +4,6 @@
 from .utils import get_split_masks
 from .utils import save_benchmark_data, write_to_csv
 from .utils import test
-from .utils import xpu_profiler
 
 __all__ = [
     'emit_itt',
@@ -15,5 +14,4 @@
     'save_benchmark_data',
     'write_to_csv',
     'test',
-    'xpu_profiler',
 ]

out.txt

@@ -0,0 +1,82 @@
+diff --git a/test/profile/test_profile.py b/test/profile/test_profile.py
+index c4b5c38f..cbd39a04 100644
+--- a/test/profile/test_profile.py
++++ b/test/profile/test_profile.py
+@@ -11,9 +11,10 @@ from torch_geometric.profile import (
+     rename_profile_file,
+     timeit,
+ )
+-from torch_geometric.profile.profile import torch_profile
++from torch_geometric.profile.profile import torch_profile, xpu_profile
+ from torch_geometric.testing import (
+     onlyCUDA,
++    onlyXPU,
+     onlyLinux,
+     onlyOnline,
+     withCUDA,
+@@ -105,3 +106,20 @@ def test_torch_profile(capfd, get_dataset, device):
+     rename_profile_file('test_profile')
+     assert os.path.exists('profile-test_profile.json')
+     os.remove('profile-test_profile.json')
++
++
++@onlyXPU
++@onlyOnline
++def test_torch_profile(capfd, get_dataset, device):
++    dataset = get_dataset(name='Cora')
++    data = dataset[0].to(device)
++    model = GraphSAGE(dataset.num_features, hidden_channels=64, num_layers=3,
++                      out_channels=dataset.num_classes).to(device)
++
++    with xpu_profile():
++        model(data.x, data.edge_index)
++
++    out, _ = capfd.readouterr()
++    assert 'Self CPU' in out
++    if data.x.is_xpu:
++        assert 'Self XPU' in out
+diff --git a/torch_geometric/testing/__init__.py b/torch_geometric/testing/__init__.py
+index 83f98204..6a108d42 100644
+--- a/torch_geometric/testing/__init__.py
++++ b/torch_geometric/testing/__init__.py
+@@ -4,6 +4,7 @@ from .decorators import (
+     onlyLinux,
+     onlyPython,
+     onlyCUDA,
++    onlyXPU,
+     onlyOnline,
+     onlyGraphviz,
+     onlyNeighborSampler,
+@@ -22,6 +23,7 @@ __all__ = [
+     'onlyLinux',
+     'onlyPython',
+     'onlyCUDA',
++    'onlyXPU',
+     'onlyOnline',
+     'onlyGraphviz',
+     'onlyNeighborSampler',
+diff --git a/torch_geometric/testing/decorators.py b/torch_geometric/testing/decorators.py
+index b62625fa..e0d61db0 100644
+--- a/torch_geometric/testing/decorators.py
++++ b/torch_geometric/testing/decorators.py
+@@ -59,6 +59,20 @@ def onlyCUDA(func: Callable) -> Callable:
+     )(func)
+
+
++def onlyXPU(func: Callable) -> Callable:
++    r"""A decorator to skip tests if XPU is not found."""
++    import pytest
++    try:
++        import intel_extension_for_pytorch
++        xpu_available = torch.xpu.is_available()
++    except ImportError:
++        xpu_available = False
++    return pytest.mark.skipif(
++        not xpu_available,
++        reason="XPU not available",
++    )(func)
++
++
+ def onlyOnline(func: Callable):
+     r"""A decorator to skip tests if there exists no connection to the
+     internet."""

torch_geometric/profile/__init__.py

@@ -4,6 +4,7 @@
     print_time_total,
     rename_profile_file,
     torch_profile,
+    xpu_profile,
 )
 from .utils import count_parameters
 from .utils import get_model_size
@@ -21,6 +22,7 @@
     'print_time_total',
     'rename_profile_file',
     'torch_profile',
+    'xpu_profile',
     'count_parameters',
     'get_model_size',
     'get_data_size',