Repository: pyg-team/pytorch_geometric
Pull request #7994 (2023-09-07 06:15:27)

Added internal `cumsum` implementation

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7994](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (1c6725e) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/be6ce568d2f175bdb6c0842cb38d7c7d343d88b2?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (be6ce56) will **decrease** coverage by `0.73%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #7994      +/-   ##
==========================================
- Coverage   90.29%   89.57%   -0.73%     
==========================================
  Files         460      461       +1     
  Lines       26965    26960       -5     
==========================================
- Hits        24349    24149     -200     
- Misses       2616     2811     +195     
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/data/collate.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvY29sbGF0ZS5weQ==) | `89.61% <100.00%> (-0.39%)` | :arrow_down: |
| [torch\_geometric/io/tu.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2lvL3R1LnB5) | `88.37% <100.00%> (-2.54%)` | :arrow_down: |
| [torch\_geometric/nn/aggr/quantile.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2FnZ3IvcXVhbnRpbGUucHk=) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/pool/decimation.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvZGVjaW1hdGlvbi5weQ==) | `93.33% <100.00%> (+0.47%)` | :arrow_up: |
| [torch\_geometric/nn/pool/select/topk.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvc2VsZWN0L3RvcGsucHk=) | `94.54% <100.00%> (-0.10%)` | :arrow_down: |
| [torch\_geometric/nn/to\_hetero\_module.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3RvX2hldGVyb19tb2R1bGUucHk=) | `91.13% <100.00%> (-1.27%)` | :arrow_down: |
| [torch\_geometric/transforms/line\_graph.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3RyYW5zZm9ybXMvbGluZV9ncmFwaC5weQ==) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/utils/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/utils/augmentation.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2F1Z21lbnRhdGlvbi5weQ==) | `95.23% <100.00%> (ø)` | |
| [torch\_geometric/utils/dropout.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2Ryb3BvdXQucHk=) | `93.75% <100.00%> (-0.30%)` | :arrow_down: |
| ... and [6 more](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | |

... and [27 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7994/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?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 f2170194a..5ddf75080 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 `utils.cumsum` implementation ([#7994](https://github.com/pyg-team/pytorch_geometric/pull/7994))
 - Added the `BrcaTcga` dataset ([#7905](https://github.com/pyg-team/pytorch_geometric/pull/7905))
 - Added the `MyketDataset` ([#7959](https://github.com/pyg-team/pytorch_geometric/pull/7959))
 - Added a multi-GPU `ogbn-papers100M` example ([#7921](https://github.com/pyg-team/pytorch_geometric/pull/7921))
diff --git a/torch_geometric/data/collate.py b/torch_geometric/data/collate.py
index f217479b6..bef83df98 100644
--- a/torch_geometric/data/collate.py
+++ b/torch_geometric/data/collate.py
@@ -1,6 +1,6 @@
 from collections import defaultdict
 from collections.abc import Mapping, Sequence
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -9,7 +9,7 @@ import torch_geometric.typing
 from torch_geometric.data.data import BaseData
 from torch_geometric.data.storage import BaseStorage, NodeStorage
 from torch_geometric.typing import SparseTensor, torch_sparse
-from torch_geometric.utils import is_sparse, is_torch_sparse_tensor
+from torch_geometric.utils import cumsum, is_sparse, is_torch_sparse_tensor
 from torch_geometric.utils.sparse import cat
 
 
@@ -131,7 +131,8 @@ def _collate(
         cat_dim = data_list[0].__cat_dim__(key, elem, stores[0])
         if cat_dim is None or elem.dim() == 0:
             values = [value.unsqueeze(0) for value in values]
-        slices = cumsum([value.size(cat_dim or 0) for value in values])
+        sizes = torch.tensor([value.size(cat_dim or 0) for value in values])
+        slices = cumsum(sizes)
         if increment:
             incs = get_incs(key, values, data_list, stores)
             if incs.dim() > 1 or int(incs[-1]) != 0:
@@ -179,7 +180,7 @@ def _collate(
         cat_dim = data_list[0].__cat_dim__(key, elem, stores[0])
         cat_dims = (cat_dim, ) if isinstance(cat_dim, int) else cat_dim
         repeats = [[value.size(dim) for dim in cat_dims] for value in values]
-        slices = cumsum(repeats)
+        slices = cumsum(torch.tensor(repeats))
         if is_torch_sparse_tensor(elem):
             value = cat(values, dim=cat_dim)
         else:
@@ -268,15 +269,6 @@ def repeat_interleave(
     return torch.cat(outs, dim=0)
 
 
-def cumsum(value: Union[Tensor, List[int]]) -> Tensor:
-    if not isinstance(value, Tensor):
-        value = torch.tensor(value)
-    out = value.new_empty((value.size(0) + 1, ) + value.size()[1:])
-    out[0] = 0
-    torch.cumsum(value, 0, out=out[1:])
-    return out
-
-
 def get_incs(key, values: List[Any], data_list: List[BaseData],
              stores: List[BaseStorage]) -> Tensor:
     repeats = [
diff --git a/torch_geometric/datasets/upfd.py b/torch_geometric/datasets/upfd.py
index 7cd84a182..fcaf39fe7 100644
--- a/torch_geometric/datasets/upfd.py
+++ b/torch_geometric/datasets/upfd.py
@@ -13,7 +13,7 @@ from torch_geometric.data import (
     extract_zip,
 )
 from torch_geometric.io import read_txt_array
-from torch_geometric.utils import coalesce
+from torch_geometric.utils import coalesce, cumsum
 
 
 class UPFD(InMemoryDataset):
@@ -139,10 +139,8 @@ class UPFD(InMemoryDataset):
         batch = np.load(osp.join(self.raw_dir, 'node_graph_id.npy'))
         batch = torch.from_numpy(batch).to(torch.long)
 
-        node_slice = torch.cumsum(batch.bincount(), 0)
-        node_slice = torch.cat([torch.tensor([0]), node_slice])
-        edge_slice = torch.cumsum(batch[edge_index[0]].bincount(), 0)
-        edge_slice = torch.cat([torch.tensor([0]), edge_slice])
+        node_slice = cumsum(batch.bincount())
+        edge_slice = cumsum(batch[edge_index[0].bincount()])
         graph_slice = torch.arange(y.size(0) + 1)
         self.slices = {
             'x': node_slice,
diff --git a/torch_geometric/io/tu.py b/torch_geometric/io/tu.py
index 0083a3806..60940924e 100644
--- a/torch_geometric/io/tu.py
+++ b/torch_geometric/io/tu.py
@@ -7,7 +7,7 @@ import torch
 
 from torch_geometric.data import Data
 from torch_geometric.io import read_txt_array
-from torch_geometric.utils import coalesce, one_hot, remove_self_loops
+from torch_geometric.utils import coalesce, cumsum, one_hot, remove_self_loops
 
 names = [
     'A', 'graph_indicator', 'node_labels', 'node_attributes'
@@ -100,12 +100,10 @@ def cat(seq):
 
 
 def split(data, batch):
-    node_slice = torch.cumsum(torch.from_numpy(np.bincount(batch)), 0)
-    node_slice = torch.cat([torch.tensor([0]), node_slice])
+    node_slice = cumsum(torch.from_numpy(np.bincount(batch)))
 
     row, _ = data.edge_index
-    edge_slice = torch.cumsum(torch.from_numpy(np.bincount(batch[row])), 0)
-    edge_slice = torch.cat([torch.tensor([0]), edge_slice])
+    edge_slice = cumsum(torch.from_numpy(np.bincount(batch[row])))
 
     # Edge indices should start at zero for every graph.
     data.edge_index -= node_slice[batch[row]].unsqueeze(0)
diff --git a/torch_geometric/nn/aggr/quantile.py b/torch_geometric/nn/aggr/quantile.py
index 6c20b2a81..859edbb5a 100644
--- a/torch_geometric/nn/aggr/quantile.py
+++ b/torch_geometric/nn/aggr/quantile.py
@@ -4,6 +4,7 @@ import torch
 from torch import Tensor
 
 from torch_geometric.nn.aggr import Aggregation
+from torch_geometric.utils import cumsum
 
 
 class QuantileAggregation(Aggregation):
@@ -76,14 +77,14 @@ class QuantileAggregation(Aggregation):
         assert index is not None  # Required for TorchScript.
 
         count = torch.bincount(index, minlength=dim_size or 0)
-        cumsum = torch.cumsum(count, dim=0) - count
+        ptr = cumsum(count)[:-1]
 
         # In case there exists dangling indices (`dim_size > index.max()`), we
         # need to clamp them to prevent out-of-bound issues:
         if dim_size is not None:
-            cumsum = cumsum.clamp(max=x.size(dim) - 1)
+            ptr = ptr.clamp(max=x.size(dim) - 1)
 
-        q_point = self.q * (count - 1) + cumsum
+        q_point = self.q * (count - 1) + ptr
         q_point = q_point.t().reshape(-1)
 
         shape = [1] * x.dim()
diff --git a/torch_geometric/nn/data_parallel.py b/torch_geometric/nn/data_parallel.py
index b3f55b3da..4956085be 100644
--- a/torch_geometric/nn/data_parallel.py
+++ b/torch_geometric/nn/data_parallel.py
@@ -4,6 +4,7 @@ from itertools import chain
 import torch
 
 from torch_geometric.data import Batch
+from torch_geometric.utils import cumsum
 
 
 class DataParallel(torch.nn.DataParallel):
@@ -74,13 +75,11 @@ class DataParallel(torch.nn.DataParallel):
         num_devices = min(len(device_ids), len(data_list))
 
         count = torch.tensor([data.num_nodes for data in data_list])
-        cumsum = count.cumsum(0)
-        cumsum = torch.cat([cumsum.new_zeros(1), cumsum], dim=0)
-        device_id = num_devices * cumsum.to(torch.float) / cumsum[-1].item()
+        ptr = cumsum(count)
+        device_id = num_devices * ptr.to(torch.float) / ptr[-1].item()
         device_id = (device_id[:-1] + device_id[1:]) / 2.0
         device_id = device_id.to(torch.long)  # round.
-        split = device_id.bincount().cumsum(0)
-        split = torch.cat([split.new_zeros(1), split], dim=0)
+        split = cumsum(device_id.bincount())
         split = torch.unique(split, sorted=True)
         split = split.tolist()
 
diff --git a/torch_geometric/nn/pool/decimation.py b/torch_geometric/nn/pool/decimation.py
index fbc1801b9..460b0031b 100644
--- a/torch_geometric/nn/pool/decimation.py
+++ b/torch_geometric/nn/pool/decimation.py
@@ -3,6 +3,8 @@ from typing import Tuple, Union
 import torch
 from torch import LongTensor, Tensor
 
+from torch_geometric.utils import cumsum
+
 
 def decimation_indices(
     ptr: LongTensor,
@@ -40,6 +42,6 @@ def decimation_indices(
     decim_indices = torch.cat(decim_indices, dim=0)
 
     # Get updated ptr (e.g., for future decimations):
-    decim_ptr = torch.cat([decim_count.new_zeros(1), decim_count.cumsum(0)])
+    decim_ptr = cumsum(decim_count)
 
     return decim_indices, decim_ptr
diff --git a/torch_geometric/nn/pool/select/topk.py b/torch_geometric/nn/pool/select/topk.py
index 76e15126e..244b64577 100644
--- a/torch_geometric/nn/pool/select/topk.py
+++ b/torch_geometric/nn/pool/select/topk.py
@@ -6,7 +6,7 @@ from torch import Tensor
 from torch_geometric.nn.inits import uniform
 from torch_geometric.nn.pool.select import Select, SelectOutput
 from torch_geometric.nn.resolver import activation_resolver
-from torch_geometric.utils import scatter, softmax
+from torch_geometric.utils import cumsum, scatter, softmax
 
 
 # TODO (matthias) Document this method.
@@ -38,8 +38,7 @@ def topk(
         batch, batch_perm = torch.sort(batch, descending=False, stable=True)
 
         arange = torch.arange(x.size(0), dtype=torch.long, device=x.device)
-        ptr = num_nodes.new_zeros(num_nodes.numel() + 1)
-        torch.cumsum(num_nodes, 0, out=ptr[1:])
+        ptr = cumsum(num_nodes)
         batched_arange = arange - ptr[batch]
         mask = batched_arange < k[batch]
 
diff --git a/torch_geometric/nn/to_hetero_module.py b/torch_geometric/nn/to_hetero_module.py
index b8c40a52f..0aa4b29c0 100644
--- a/torch_geometric/nn/to_hetero_module.py
+++ b/torch_geometric/nn/to_hetero_module.py
@@ -8,7 +8,7 @@ from torch import Tensor
 
 import torch_geometric
 from torch_geometric.typing import EdgeType, NodeType, OptTensor
-from torch_geometric.utils import scatter
+from torch_geometric.utils import cumsum, scatter
 
 
 class ToHeteroLinear(torch.nn.Module):
@@ -126,7 +126,7 @@ class ToHeteroMessagePassing(torch.nn.Module):
         edge_sizes = scatter(torch.ones_like(edge_type), edge_type, dim=0,
                              dim_size=len(self.edge_types), reduce='sum')
 
-        cumsum = torch.cat([node_type.new_zeros(1), node_sizes.cumsum(0)[:1]])
+        ptr = cumsum(node_sizes)
 
         xs = x.split(node_sizes.tolist())
         x_dict = {node_type: x for node_type, x in zip(self.node_types, xs)}
@@ -134,8 +134,8 @@ class ToHeteroMessagePassing(torch.nn.Module):
         # TODO Consider out-sourcing to its own function.
         edge_indices = edge_index.clone().split(edge_sizes.tolist(), dim=1)
         for (src, _, dst), index in zip(self.edge_types, edge_indices):
-            index[0] -= cumsum[self.node_type_to_index[src]]
-            index[1] -= cumsum[self.node_type_to_index[dst]]
+            index[0] -= ptr[self.node_type_to_index[src]]
+            index[1] -= ptr[self.node_type_to_index[dst]]
 
         edge_index_dict = {
             edge_type: edge_index
diff --git a/torch_geometric/transforms/line_graph.py b/torch_geometric/transforms/line_graph.py
index 83c492b26..ae2c5c1bf 100644
--- a/torch_geometric/transforms/line_graph.py
+++ b/torch_geometric/transforms/line_graph.py
@@ -3,7 +3,7 @@ import torch
 from torch_geometric.data import Data
 from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
-from torch_geometric.utils import coalesce, remove_self_loops, scatter
+from torch_geometric.utils import coalesce, cumsum, remove_self_loops, scatter
 
 
 @functional_transform('line_graph')
@@ -45,12 +45,9 @@ class LineGraph(BaseTransform):
 
             count = scatter(torch.ones_like(row), row, dim=0,
                             dim_size=data.num_nodes, reduce='sum')
-            cumsum = torch.cat([count.new_zeros(1), count.cumsum(0)], dim=0)
+            ptr = cumsum(count)
 
-            cols = [
-                i[cumsum[col[j]]:cumsum[col[j] + 1]]
-                for j in range(col.size(0))
-            ]
+            cols = [i[ptr[col[j]]:ptr[col[j] + 1]] for j in range(col.size(0))]
             rows = [row.new_full((c.numel(), ), j) for j, c in enumerate(cols)]
 
             row, col = torch.cat(rows, dim=0), torch.cat(cols, dim=0)
diff --git a/torch_geometric/utils/__init__.py b/torch_geometric/utils/__init__.py
index 30c648a61..faa52fea5 100644
--- a/torch_geometric/utils/__init__.py
+++ b/torch_geometric/utils/__init__.py
@@ -3,9 +3,9 @@ import copy
 from .scatter import scatter, group_argsort
 from .segment import segment
 from .sort import index_sort
+from .functions import cumsum
 from .degree import degree
 from .softmax import softmax
-from .dropout import dropout_adj, dropout_node, dropout_edge, dropout_path
 from .sort_edge_index import sort_edge_index
 from .lexsort import lexsort
 from .coalesce import coalesce
@@ -16,6 +16,7 @@ from .loop import (contains_self_loops, remove_self_loops,
 from .isolated import contains_isolated_nodes, remove_isolated_nodes
 from .subgraph import (get_num_hops, subgraph, k_hop_subgraph,
                        bipartite_subgraph)
+from .dropout import dropout_adj, dropout_node, dropout_edge, dropout_path
 from .homophily import homophily
 from .assortativity import assortativity
 from .get_laplacian import get_laplacian
@@ -59,12 +60,9 @@ __all__ = [
     'group_argsort',
     'segment',
     'index_sort',
+    'cumsum',
     'degree',
     'softmax',
-    'dropout_node',
-    'dropout_edge',
-    'dropout_path',
-    'dropout_adj',
     'sort_edge_index',
     'lexsort',
     'coalesce',
@@ -82,6 +80,10 @@ __all__ = [
     'subgraph',
     'bipartite_subgraph',
     'k_hop_subgraph',
+    'dropout_node',
+    'dropout_edge',
+    'dropout_path',
+    'dropout_adj',
     'homophily',
     'assortativity',
     'get_laplacian',
diff --git a/torch_geometric/utils/augmentation.py b/torch_geometric/utils/augmentation.py
index 830145319..59e66824a 100644
--- a/torch_geometric/utils/augmentation.py
+++ b/torch_geometric/utils/augmentation.py
@@ -3,7 +3,7 @@ from typing import Optional, Tuple, Union
 import torch
 from torch import Tensor
 
-from torch_geometric.utils import negative_sampling, scatter
+from torch_geometric.utils import cumsum, negative_sampling, scatter
 
 
 def shuffle_node(
@@ -61,10 +61,10 @@ def shuffle_node(
         perm = torch.randperm(x.size(0), device=x.device)
         return x[perm], perm
     num_nodes = scatter(batch.new_ones(x.size(0)), batch, dim=0, reduce='sum')
-    cumsum = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])
+    ptr = cumsum(num_nodes)
     perm = torch.cat([
         torch.randperm(n, device=x.device) + offset
-        for offset, n in zip(cumsum[:-1], num_nodes)
+        for offset, n in zip(ptr[:-1], num_nodes)
     ])
     return x[perm], perm
 
diff --git a/torch_geometric/utils/dropout.py b/torch_geometric/utils/dropout.py
index 9a90953c1..c85f8bf94 100644
--- a/torch_geometric/utils/dropout.py
+++ b/torch_geometric/utils/dropout.py
@@ -6,10 +6,8 @@ from torch import Tensor
 import torch_geometric.typing
 from torch_geometric.deprecation import deprecated
 from torch_geometric.typing import OptTensor
-from torch_geometric.utils.degree import degree
+from torch_geometric.utils import cumsum, degree, sort_edge_index, subgraph
 from torch_geometric.utils.num_nodes import maybe_num_nodes
-from torch_geometric.utils.sort_edge_index import sort_edge_index
-from torch_geometric.utils.subgraph import subgraph
 
 
 def filter_adj(row: Tensor, col: Tensor, edge_attr: OptTensor,
@@ -277,15 +275,12 @@ def dropout_path(edge_index: Tensor, p: float = 0.2, walks_per_node: int = 1,
     sample_mask = torch.rand(row.size(0), device=edge_index.device) <= p
     start = row[sample_mask].repeat(walks_per_node)
 
-    deg = degree(row, num_nodes=num_nodes)
-    rowptr = row.new_zeros(num_nodes + 1)
-    torch.cumsum(deg, 0, out=rowptr[1:])
+    rowptr = cumsum(degree(row, num_nodes=num_nodes, dtype=torch.long))
     n_id, e_id = torch.ops.torch_cluster.random_walk(rowptr, col, start,
                                                      walk_length, 1.0, 1.0)
     e_id = e_id[e_id != -1].view(-1)  # filter illegal edges
 
-    if edge_orders is not None:
-        # permute edge ids
+    if edge_orders is not None:  # Permute edge indices:
         e_id = edge_orders[e_id]
     edge_mask[e_id] = False
     edge_index = ori_edge_index[:, edge_mask]
diff --git a/torch_geometric/utils/functions.py b/torch_geometric/utils/functions.py
new file mode 100644
index 000000000..6708de7fe
--- /dev/null
+++ b/torch_geometric/utils/functions.py
@@ -0,0 +1,27 @@
+import torch
+from torch import Tensor
+
+
+def cumsum(x: Tensor, dim: int = 0) -> Tensor:
+    r"""Returns the cumulative sum of elements of :obj:`inputs`.
+    In contrast to :meth:`torch.cumsum`, preprends the output with a zero.
+
+    Args:
+        x (torch.Tensor): The input tensor.
+        dim (int, optional): The dimension to do the operation over.
+            (default: :obj:`0`)
+
+    Example:
+
+        >>> x = torch.tensor([2, 4, 1])
+        >>> cumsum(x)
+        tensor([0, 2, 6, 7])
+
+    """
+    size = x.size()[:dim] + (x.size(dim) + 1, ) + x.size()[dim + 1:]
+    out = x.new_empty(size)
+
+    out.narrow(dim, 0, 1).zero_()
+    torch.cumsum(x, dim=dim, out=out.narrow(dim, 1, x.size(dim)))
+
+    return out
diff --git a/torch_geometric/utils/negative_sampling.py b/torch_geometric/utils/negative_sampling.py
index b7e0d0f09..134feced4 100644
--- a/torch_geometric/utils/negative_sampling.py
+++ b/torch_geometric/utils/negative_sampling.py
@@ -5,7 +5,7 @@ import numpy as np
 import torch
 from torch import Tensor
 
-from torch_geometric.utils import coalesce, degree, remove_self_loops
+from torch_geometric.utils import coalesce, cumsum, degree, remove_self_loops
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 
 
@@ -178,25 +178,25 @@ def batched_negative_sampling(
     edge_indices = torch.split(edge_index, split, dim=1)
 
     num_src = degree(src_batch, dtype=torch.long)
-    cum_src = torch.cat([src_batch.new_zeros(1), num_src.cumsum(0)[:-1]])
+    cum_src = cumsum(num_src)[:-1]
 
     if isinstance(batch, Tensor):
         num_nodes = num_src.tolist()
-        cumsum = cum_src
+        ptr = cum_src
     else:
         num_dst = degree(dst_batch, dtype=torch.long)
-        cum_dst = torch.cat([dst_batch.new_zeros(1), num_dst.cumsum(0)[:-1]])
+        cum_dst = cumsum(num_dst)[:-1]
 
         num_nodes = torch.stack([num_src, num_dst], dim=1).tolist()
-        cumsum = torch.stack([cum_src, cum_dst], dim=1).unsqueeze(-1)
+        ptr = torch.stack([cum_src, cum_dst], dim=1).unsqueeze(-1)
 
     neg_edge_indices = []
     for i, edge_index in enumerate(edge_indices):
-        edge_index = edge_index - cumsum[i]
+        edge_index = edge_index - ptr[i]
         neg_edge_index = negative_sampling(edge_index, num_nodes[i],
                                            num_neg_samples, method,
                                            force_undirected)
-        neg_edge_index += cumsum[i]
+        neg_edge_index += ptr[i]
         neg_edge_indices.append(neg_edge_index)
 
     return torch.cat(neg_edge_indices, dim=1)
diff --git a/torch_geometric/utils/scatter.py b/torch_geometric/utils/scatter.py
index b86f96d4c..d9207c90a 100644
--- a/torch_geometric/utils/scatter.py
+++ b/torch_geometric/utils/scatter.py
@@ -6,6 +6,7 @@ from torch import Tensor
 import torch_geometric.typing
 from torch_geometric import warnings
 from torch_geometric.typing import torch_scatter
+from torch_geometric.utils.functions import cumsum
 
 if torch_geometric.typing.WITH_PT112:  # pragma: no cover
 
@@ -250,7 +251,6 @@ def group_argsort(
     # Compute cumulative sum of number of entries with the same index:
     count = scatter(torch.ones_like(index), index, dim=dim,
                     dim_size=num_groups, reduce='sum')
-    ptr = count.new_zeros(count.numel() + 1)
-    torch.cumsum(count, dim=0, out=ptr[1:])
+    ptr = cumsum(count)
 
     return out - ptr[index]
diff --git a/torch_geometric/utils/to_dense_adj.py b/torch_geometric/utils/to_dense_adj.py
index e2423735a..1a501b40b 100644
--- a/torch_geometric/utils/to_dense_adj.py
+++ b/torch_geometric/utils/to_dense_adj.py
@@ -4,7 +4,7 @@ import torch
 from torch import Tensor
 
 from torch_geometric.typing import OptTensor
-from torch_geometric.utils import scatter
+from torch_geometric.utils import cumsum, scatter
 
 
 def to_dense_adj(
@@ -67,7 +67,7 @@ def to_dense_adj(
 
     one = batch.new_ones(batch.size(0))
     num_nodes = scatter(one, batch, dim=0, dim_size=batch_size, reduce='sum')
-    cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])
+    cum_nodes = cumsum(num_nodes)
 
     idx0 = batch[edge_index[0]]
     idx1 = edge_index[0] - cum_nodes[batch][edge_index[0]]
diff --git a/torch_geometric/utils/to_dense_batch.py b/torch_geometric/utils/to_dense_batch.py
index cad023e7b..836f3b884 100644
--- a/torch_geometric/utils/to_dense_batch.py
+++ b/torch_geometric/utils/to_dense_batch.py
@@ -7,7 +7,7 @@ from torch_geometric.experimental import (
     disable_dynamic_shapes,
     is_experimental_mode_enabled,
 )
-from torch_geometric.utils import scatter
+from torch_geometric.utils import cumsum, scatter
 
 
 @disable_dynamic_shapes(required_args=['batch_size', 'max_num_nodes'])
@@ -106,7 +106,7 @@ def to_dense_batch(
 
     num_nodes = scatter(batch.new_ones(x.size(0)), batch, dim=0,
                         dim_size=batch_size, reduce='sum')
-    cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])
+    cum_nodes = cumsum(num_nodes)
 
     filter_nodes = False
     dynamic_shapes_disabled = is_experimental_mode_enabled(
diff --git a/torch_geometric/utils/unbatch.py b/torch_geometric/utils/unbatch.py
index 0ba85c7e1..304f60d2b 100644
--- a/torch_geometric/utils/unbatch.py
+++ b/torch_geometric/utils/unbatch.py
@@ -3,7 +3,7 @@ from typing import List, Optional
 import torch
 from torch import Tensor
 
-from torch_geometric.utils import degree
+from torch_geometric.utils import cumsum, degree
 
 
 def unbatch(
@@ -65,7 +65,7 @@ def unbatch_edge_index(
                 [1, 0, 2, 1]]))
     """
     deg = degree(batch, batch_size, dtype=torch.long)
-    ptr = torch.cat([deg.new_zeros(1), deg.cumsum(dim=0)[:-1]], dim=0)
+    ptr = cumsum(deg)
 
     edge_batch = batch[edge_index[0]]
     edge_index = edge_index - ptr[edge_batch]


Test patch:
diff --git a/test/nn/dense/test_linear.py b/test/nn/dense/test_linear.py
index 9c004b7b0..f1189286e 100644
--- a/test/nn/dense/test_linear.py
+++ b/test/nn/dense/test_linear.py
@@ -12,6 +12,7 @@ from torch_geometric.nn import HeteroDictLinear, HeteroLinear, Linear
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import withCUDA, withPackage
 from torch_geometric.typing import pyg_lib
+from torch_geometric.utils import cumsum
 
 weight_inits = ['glorot', 'kaiming_uniform', None]
 bias_inits = ['zeros', None]
@@ -300,7 +301,7 @@ if __name__ == '__main__':
 
         xs = get_xs(mean, std, num_types, channels)
         count = torch.tensor([x.size(0) for x in xs])
-        ptr = torch.tensor([0] + [x.size(0) for x in xs]).cumsum(0)
+        ptr = cumsum(torch.tensor([x.size(0) for x in xs]))
         x = torch.cat(xs, dim=0)
         padded_x = torch.nested.nested_tensor(xs).to_padded_tensor(padding=0.0)
         weight = torch.randn(num_types, channels, channels, device=args.device)
diff --git a/test/utils/test_functions.py b/test/utils/test_functions.py
new file mode 100644
index 000000000..129387270
--- /dev/null
+++ b/test/utils/test_functions.py
@@ -0,0 +1,11 @@
+import torch
+
+from torch_geometric.utils import cumsum
+
+
+def test_cumsum():
+    x = torch.tensor([2, 4, 1])
+    assert cumsum(x).tolist() == [0, 2, 6, 7]
+
+    x = torch.tensor([[2, 4], [3, 6]])
+    assert cumsum(x, dim=1).tolist() == [[0, 2, 6], [0, 3, 9]]