Repository: pyg-team/pytorch_geometric
Pull request #8399 (2023-11-17 13:32:02)

Use optimized implementation in softmax operation
This PR uses optimized `softmax_csr` operation (introduced in [pyg-lib @ 264](https://github.com/pyg-team/pyg-lib/pull/264)), when given is a CPU tensor, and softmax groups are defined via `ptr`.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8399?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `2 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`cf24b4b`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/cf24b4bcb4e825537ba08d8fc5f31073e2cd84c7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.73% compared to head [(`263a461`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8399?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.73%.

> :exclamation: Current head 263a461 differs from pull request most recent head b466487. Consider uploading reports for the commit b466487 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8399?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/typing.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8399?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3R5cGluZy5weQ==) | 50.00% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8399?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |
| [torch\_geometric/utils/softmax.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8399?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3NvZnRtYXgucHk=) | 75.00% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8399?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    #8399      +/-   ##
==========================================
- Coverage   87.73%   87.73%   -0.01%     
==========================================
  Files         475      475              
  Lines       28890    28896       +6     
==========================================
+ Hits        25348    25352       +4     
- Misses       3542     3544       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8399?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 a305341f2..54986feac 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Made `utils.softmax` faster via `softmax_csr` ([#8399](https://github.com/pyg-team/pytorch_geometric/pull/8399))
 - Made `utils.mask.mask_select` faster ([#8369](https://github.com/pyg-team/pytorch_geometric/pull/8369))
 - Update `DistNeighborSampler` for homogeneous graphs ([#8209](https://github.com/pyg-team/pytorch_geometric/pull/8209), [#8367](https://github.com/pyg-team/pytorch_geometric/pull/8367))
 - Update `GraphStore` and `FeatureStore` to support distributed training ([#8083](https://github.com/pyg-team/pytorch_geometric/pull/8083))
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index 085ecc613..38572910e 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -40,6 +40,7 @@ try:
             WITH_GMM = False
             WITH_SEGMM = False
     WITH_SAMPLED_OP = hasattr(pyg_lib.ops, 'sampled_add')
+    WITH_SOFTMAX = hasattr(pyg_lib.ops, 'softmax_csr')
     WITH_INDEX_SORT = hasattr(pyg_lib.ops, 'index_sort')
     WITH_METIS = hasattr(pyg_lib, 'partition')
     WITH_EDGE_TIME_NEIGHBOR_SAMPLE = ('edge_time' in inspect.signature(
@@ -55,6 +56,7 @@ except Exception as e:
     WITH_GMM = False
     WITH_SEGMM = False
     WITH_SAMPLED_OP = False
+    WITH_SOFTMAX = False
     WITH_INDEX_SORT = False
     WITH_METIS = False
     WITH_EDGE_TIME_NEIGHBOR_SAMPLE = False
diff --git a/torch_geometric/utils/softmax.py b/torch_geometric/utils/softmax.py
index 8003f9de6..e8f5c3c42 100644
--- a/torch_geometric/utils/softmax.py
+++ b/torch_geometric/utils/softmax.py
@@ -2,6 +2,8 @@ from typing import Optional
 
 from torch import Tensor
 
+import torch_geometric.typing
+from torch_geometric.typing import pyg_lib
 from torch_geometric.utils import scatter, segment
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 
@@ -50,6 +52,9 @@ def softmax(
                 [0.8062, 0.1938, 1.0000, 1.0000]])
     """
     if ptr is not None:
+        if (src.device.type == 'cpu'
+                and torch_geometric.typing.WITH_SOFTMAX):  # pragma: no cover
+            return pyg_lib.ops.softmax_csr(src, ptr, dim)
         dim = dim + src.dim() if dim < 0 else dim
         size = ([1] * dim) + [-1]
         count = ptr[1:] - ptr[:-1]


Test patch:
diff --git a/test/utils/test_softmax.py b/test/utils/test_softmax.py
index bfffba30c..731a389bc 100644
--- a/test/utils/test_softmax.py
+++ b/test/utils/test_softmax.py
@@ -5,6 +5,9 @@ import torch_geometric.typing
 from torch_geometric.profile import benchmark
 from torch_geometric.utils import softmax
 
+CALCULATION_VIA_PTR_AVAILABLE = (torch_geometric.typing.WITH_SOFTMAX
+                                 or torch_geometric.typing.WITH_TORCH_SCATTER)
+
 
 def test_softmax():
     src = torch.tensor([1., 1., 1., 1.])
@@ -13,7 +16,7 @@ def test_softmax():
 
     out = softmax(src, index)
     assert out.tolist() == [0.5, 0.5, 1, 1]
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert softmax(src, None, ptr).tolist() == out.tolist()
     else:
         with pytest.raises(ImportError):
@@ -22,7 +25,7 @@ def test_softmax():
     src = src.view(-1, 1)
     out = softmax(src, index)
     assert out.tolist() == [[0.5], [0.5], [1], [1]]
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert softmax(src, None, ptr).tolist() == out.tolist()
 
     jit = torch.jit.script(softmax)
@@ -52,22 +55,22 @@ def test_softmax_dim():
 
     src = torch.randn(4)
     assert torch.allclose(softmax(src, index, dim=0), src.softmax(dim=0))
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert torch.allclose(softmax(src, ptr=ptr, dim=0), src.softmax(dim=0))
 
     src = torch.randn(4, 16)
     assert torch.allclose(softmax(src, index, dim=0), src.softmax(dim=0))
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert torch.allclose(softmax(src, ptr=ptr, dim=0), src.softmax(dim=0))
 
     src = torch.randn(4, 4)
     assert torch.allclose(softmax(src, index, dim=-1), src.softmax(dim=-1))
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert torch.allclose(softmax(src, ptr=ptr, dim=-1), src.softmax(-1))
 
     src = torch.randn(4, 4, 16)
     assert torch.allclose(softmax(src, index, dim=1), src.softmax(dim=1))
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert torch.allclose(softmax(src, ptr=ptr, dim=1), src.softmax(dim=1))