[Type Hints] utils.normalized_cut (#5665)

Co-authored-by: Matthias Fey <matthias.fey@tu-dortmund.de>

CHANGELOG.md

@@ -40,7 +40,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Support `in_channels` with `tuple` in `GENConv` for bipartite message passing ([#5627](https://github.com/pyg-team/pytorch_geometric/pull/5627), [#5641](https://github.com/pyg-team/pytorch_geometric/pull/5641))
 - Handle cases of not having enough possible negative edges in `RandomLinkSplit` ([#5642](https://github.com/pyg-team/pytorch_geometric/pull/5642))
 - Fix `RGCN+pyg-lib` for `LongTensor` input ([#5610](https://github.com/pyg-team/pytorch_geometric/pull/5610))
-- Improved type hint support ([#5603](https://github.com/pyg-team/pytorch_geometric/pull/5603), [#5659](https://github.com/pyg-team/pytorch_geometric/pull/5659))
+- Improved type hint support ([#5603](https://github.com/pyg-team/pytorch_geometric/pull/5603), [#5659](https://github.com/pyg-team/pytorch_geometric/pull/5659), [#5665](https://github.com/pyg-team/pytorch_geometric/pull/5665))
 - Avoid modifying `mode_kwargs` in `MultiAggregation` ([#5601](https://github.com/pyg-team/pytorch_geometric/pull/5601))
 - Changed `BatchNorm` to allow for batches of size one during training ([#5530](https://github.com/pyg-team/pytorch_geometric/pull/5530), [#5614](https://github.com/pyg-team/pytorch_geometric/pull/5614))
 - Integrated better temporal sampling support by requiring that local neighborhoods are sorted according to time ([#5516](https://github.com/pyg-team/pytorch_geometric/issues/5516), [#5602](https://github.com/pyg-team/pytorch_geometric/issues/5602))

test/utils/test_normalized_cut.py

@@ -1,5 +1,6 @@
 import torch
 
+from torch_geometric.testing import is_full_test
 from torch_geometric.utils import normalized_cut
 
 
@@ -11,3 +12,8 @@ def test_normalized_cut():
 
     output = normalized_cut(torch.stack([row, col], dim=0), edge_attr)
     assert output.tolist() == expected_output
+
+    if is_full_test():
+        jit = torch.jit.script(normalized_cut)
+        output = jit(torch.stack([row, col], dim=0), edge_attr)
+        assert output.tolist() == expected_output

torch_geometric/utils/normalized_cut.py

@@ -1,9 +1,15 @@
 from typing import Optional
 
+from torch import Tensor
+
 from torch_geometric.utils import degree
 
 
-def normalized_cut(edge_index, edge_attr, num_nodes: Optional[int] = None):
+def normalized_cut(
+    edge_index: Tensor,
+    edge_attr: Tensor,
+    num_nodes: Optional[int] = None,
+) -> Tensor:
     r"""Computes the normalized cut :math:`\mathbf{e}_{i,j} \cdot
     \left( \frac{1}{\deg(i)} + \frac{1}{\deg(j)} \right)` of a weighted graph
     given by edge indices and edge attributes.