Repository: pyg-team/pytorch_geometric
Pull request #8852 (2024-02-02 09:29:27)

Fallback to `torch._segment_reduce`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8852?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `5 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`a6ddfc3`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/a6ddfc3f7cfc84f25dd7fb4817fa8683ad35c73b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.40% compared to head [(`4dc3ce9`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8852?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.39%.

> :exclamation: Current head 4dc3ce9 differs from pull request most recent head 0080fb5. Consider uploading reports for the commit 0080fb5 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8852?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/utils/\_segment.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8852?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL19zZWdtZW50LnB5) | 37.50% | [5 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8852?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    #8852      +/-   ##
==========================================
- Coverage   89.40%   89.39%   -0.02%     
==========================================
  Files         479      479              
  Lines       31165    31172       +7     
==========================================
+ Hits        27863    27865       +2     
- Misses       3302     3307       +5     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8852?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/CHANGELOG.md b/CHANGELOG.md
index 986858d49..d6f868deb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added fallback code path for `segment` in case `torch-scatter` is not installed ([#8852](https://github.com/pyg-team/pytorch_geometric/pull/8852))
 - Added support for custom node labels in `visualize_graph()` ([#8816](https://github.com/pyg-team/pytorch_geometric/pull/8816))
 - Added support for graph partitioning for temporal data in `torch_geometric.distributed` ([#8718](https://github.com/pyg-team/pytorch_geometric/pull/8718), [#8815](https://github.com/pyg-team/pytorch_geometric/pull/8815))
 - Added `TreeGraph` and `GridMotif` generators ([#8736](https://github.com/pyg-team/pytorch_geometric/pull/8736))
diff --git a/torch_geometric/utils/_segment.py b/torch_geometric/utils/_segment.py
index c730c94e9..786714253 100644
--- a/torch_geometric/utils/_segment.py
+++ b/torch_geometric/utils/_segment.py
@@ -1,3 +1,4 @@
+import torch
 from torch import Tensor
 
 import torch_geometric.typing
@@ -18,9 +19,28 @@ def segment(src: Tensor, ptr: Tensor, reduce: str = 'sum') -> Tensor:
             refers to the boundaries of segments such that :obj:`ptr[0] = 0`
             and :obj:`ptr[-1] = src.size(0)`.
         reduce (str, optional): The reduce operation (:obj:`"sum"`,
-            :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`).
+            :obj:`"mean"`, :obj:`"min"` or :obj:`"max"`).
             (default: :obj:`"sum"`)
     """
     if not torch_geometric.typing.WITH_TORCH_SCATTER or is_compiling():
-        raise ImportError("'segment' requires the 'torch-scatter' package")
+        return _torch_segment(src, ptr, reduce)
+
+    if torch_geometric.typing.WITH_PT20 and src.is_cuda and reduce == 'mean':
+        return _torch_segment(src, ptr, reduce)
+
+    # TODO Fallback to `scatter` if deterministic algorithms are turned off.
+
     return torch_scatter.segment_csr(src, ptr, reduce=reduce)
+
+
+def _torch_segment(src: Tensor, ptr: Tensor, reduce: str = 'sum') -> Tensor:
+    if not torch_geometric.typing.WITH_PT20:
+        raise ImportError("'segment' requires the 'torch-scatter' package")
+
+    if reduce == 'min' or reduce == 'max':
+        reduce = f'a{reduce}'  # `amin` or `amax`
+    initial = 0 if reduce == 'mean' else None
+    out = torch._segment_reduce(src, reduce, offsets=ptr, initial=initial)
+    if reduce == 'amin' or reduce == 'amax':
+        out = torch.where(out.isinf(), 0, out)
+    return out


Test patch:
diff --git a/test/utils/test_segment.py b/test/utils/test_segment.py
index 44eeb0a2c..2563833ef 100644
--- a/test/utils/test_segment.py
+++ b/test/utils/test_segment.py
@@ -3,26 +3,31 @@ from itertools import product
 import pytest
 import torch
 
+import torch_geometric.typing
 from torch_geometric.profile import benchmark
-from torch_geometric.testing import withCUDA, withPackage
+from torch_geometric.testing import withCUDA, withoutExtensions
 from torch_geometric.utils import scatter, segment
 from torch_geometric.utils.sparse import index2ptr
 
 
 @withCUDA
-@withPackage('torch_scatter')
+@withoutExtensions
 @pytest.mark.parametrize('reduce', ['sum', 'mean', 'min', 'max'])
-def test_segment(device, reduce):
+def test_segment(device, without_extensions, reduce):
     src = torch.randn(20, 16, device=device)
     ptr = torch.tensor([0, 0, 5, 10, 15, 20], device=device)
 
-    out = segment(src, ptr, reduce=reduce)
+    if without_extensions and not torch_geometric.typing.WITH_PT20:
+        with pytest.raises(ImportError, match="requires the 'torch-scatter'"):
+            segment(src, ptr, reduce=reduce)
+    else:
+        out = segment(src, ptr, reduce=reduce)
 
-    expected = getattr(torch, reduce)(src.view(4, 5, -1), dim=1)
-    expected = expected[0] if isinstance(expected, tuple) else expected
+        expected = getattr(torch, reduce)(src.view(4, 5, -1), dim=1)
+        expected = expected[0] if isinstance(expected, tuple) else expected
 
-    assert torch.allclose(out[:1], torch.zeros(1, 16, device=device))
-    assert torch.allclose(out[1:], expected)
+        assert torch.allclose(out[:1], torch.zeros(1, 16, device=device))
+        assert torch.allclose(out[1:], expected)
 
 
 if __name__ == '__main__':
@@ -48,15 +53,13 @@ if __name__ == '__main__':
     num_nodes_list = [4_000, 8_000, 16_000, 32_000, 64_000]
 
     if args.aggr == 'all':
-        aggrs = ['sum', 'mean', 'min', 'max', 'mul']
+        aggrs = ['sum', 'mean', 'min', 'max']
     else:
         aggrs = args.aggr.split(',')
 
     def pytorch_segment(x, ptr, reduce):
         if reduce == 'min' or reduce == 'max':
             reduce = f'a{aggr}'  # `amin` or `amax`
-        elif reduce == 'mul':
-            reduce = 'prod'
         return torch._segment_reduce(x, reduce, offsets=ptr)
 
     def own_segment(x, ptr, reduce):