Repository: pyg-team/pytorch_geometric
Pull request #9398 (2024-06-06 07:13:48)

Added new XPU support to the `withDevice` API
Added new XPU support to the `withDevice` API

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9398?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: Patch coverage is `33.33333%` with `6 lines` in your changes missing coverage. Please review.
> Project coverage is 86.85%. Comparing base [(`476e768`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/476e768593d91c1622b02440e25bbc93fc7f5294?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) to head [(`5e93ec4`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/5e93ec4ab602f9a21fe2bbaaf6102de778389cd3?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).
> Report is 3 commits behind head on master.

> :exclamation: **Current head 5e93ec4 differs from pull request most recent head a9fc83c**
> 
> Please [upload](https://docs.codecov.com/docs/codecov-uploader) reports for the commit a9fc83c to get more accurate results.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9398?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/testing/decorators.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9398?src=pr&el=tree&filepath=torch_geometric%2Ftesting%2Fdecorators.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3Rlc3RpbmcvZGVjb3JhdG9ycy5weQ==) | 33.33% | [6 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9398?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #9398      +/-   ##
==========================================
- Coverage   87.99%   86.85%   -1.15%     
==========================================
  Files         473      462      -11     
  Lines       30820    30512     -308     
==========================================
- Hits        27120    26500     -620     
- Misses       3700     4012     +312     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9398?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/torch_geometric/nn/norm/batch_norm.py b/torch_geometric/nn/norm/batch_norm.py
index 2fcc6e6ce..3830f5499 100644
--- a/torch_geometric/nn/norm/batch_norm.py
+++ b/torch_geometric/nn/norm/batch_norm.py
@@ -88,7 +88,7 @@ class BatchNorm(torch.nn.Module):
         return self.module(x)
 
     def __repr__(self):
-        return f'{self.__class__.__name__}({self.module.num_features})'
+        return f'{self.__class__.__name__}({self.module.extra_repr()})'
 
 
 class HeteroBatchNorm(torch.nn.Module):


Test patch:
diff --git a/test/nn/norm/test_batch_norm.py b/test/nn/norm/test_batch_norm.py
index 9846bc8f4..c548c0fa6 100644
--- a/test/nn/norm/test_batch_norm.py
+++ b/test/nn/norm/test_batch_norm.py
@@ -13,7 +13,8 @@ def test_batch_norm(device, conf):
     norm = BatchNorm(16, affine=conf, track_running_stats=conf).to(device)
     norm.reset_running_stats()
     norm.reset_parameters()
-    assert str(norm) == 'BatchNorm(16)'
+    assert str(norm) == (f'BatchNorm(16, eps=1e-05, momentum=0.1, '
+                         f'affine={conf}, track_running_stats={conf})')
 
     if is_full_test():
         torch.jit.script(norm)
diff --git a/test/profile/test_profile.py b/test/profile/test_profile.py
index d6245f2d7..082afb315 100644
--- a/test/profile/test_profile.py
+++ b/test/profile/test_profile.py
@@ -92,9 +92,10 @@ def test_profileit_xpu(get_dataset):
     warnings.filterwarnings('ignore', '.*arguments of DataFrame.drop.*')
 
     dataset = get_dataset(name='Cora')
-    data = dataset[0].cuda()
+    device = torch.device('xpu')
+    data = dataset[0].to(device)
     model = GraphSAGE(dataset.num_features, hidden_channels=64, num_layers=3,
-                      out_channels=dataset.num_classes).cuda()
+                      out_channels=dataset.num_classes).to(device)
     optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
 
     @profileit('xpu')
diff --git a/torch_geometric/testing/decorators.py b/torch_geometric/testing/decorators.py
index 029fef5a7..ca185f05a 100644
--- a/torch_geometric/testing/decorators.py
+++ b/torch_geometric/testing/decorators.py
@@ -234,6 +234,18 @@ def withDevice(func: Callable) -> Callable:
         except RuntimeError:
             pass
 
+    if not hasattr(torch, 'xpu'):
+        try:
+            import intel_extension_for_pytorch as ipex
+            xpu_available = ipex.xpu.is_available()
+        except ImportError:
+            xpu_available = False
+    else:
+        xpu_available = torch.xpu.is_available()
+
+    if xpu_available:
+        devices.append(pytest.param(torch.device('xpu:0'), id='xpu'))
+
     # Additional devices can be registered through environment variables:
     device = os.getenv('TORCH_DEVICE')
     if device: