[Type Hints] utils.grid (#5724)

This is a draft PR, as the test in `test_grid.py` is failing because of
issues with the `coalesce` function.

Currently, the code uses `torch.sparse.coalesce` , which does not have
any type hints(hence cant be used with jit). On @rusty1s suggestion, I
tried using `torch_geometric.utils.coalesce` instead.

The test (on running `FULL_TEST=1 pytest test/utils/test_grid.py`) is
still failing, with the traceback as below :
```
E           RuntimeError: 
E           'Union[Tensor, List[Tensor], NoneType]' object is not subscriptable:
E             File "/home/nihal/opensource/pytorch_geometric/torch_geometric/utils/coalesce.py", line 78
E                   edge_index = edge_index[:, perm]
E                   if edge_attr is not None and isinstance(edge_attr, Tensor):
E                       edge_attr = edge_attr[perm]
E                                   ~~~~~~~~~~~~~~~ <--- HERE
E                   elif edge_attr is not None:
E                       edge_attr = [e[perm] for e in edge_attr]
E           'coalesce' is being compiled since it was called from 'grid_index'
E             File "/home/nihal/opensource/pytorch_geometric/torch_geometric/utils/grid.py", line 60
E           
E               edge_index = torch.stack([row, col], dim=0)
E               edge_index = coalesce(edge_index, None, height * width)
E               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
E               print('------------------------',type(edge_index))
E               return edge_index
E           'grid_index' is being compiled since it was called from 'grid'
E             File "/home/nihal/opensource/pytorch_geometric/torch_geometric/utils/grid.py", line 38
E               """
E           
E               edge_index = grid_index(height, width, device)
E               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
E               pos = grid_pos(height, width, dtype, device)
E               return edge_index, pos

../../miniconda3/envs/pygenv/lib/python3.8/site-packages/torch/jit/_script.py:1343: RuntimeError

```

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

Author: nihal-rao

CHANGELOG.md

@@ -46,7 +46,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), [#5664](https://github.com/pyg-team/pytorch_geometric/pull/5664), [#5665](https://github.com/pyg-team/pytorch_geometric/pull/5665), [#5666](https://github.com/pyg-team/pytorch_geometric/pull/5666), [#5667](https://github.com/pyg-team/pytorch_geometric/pull/5667), [#5668](https://github.com/pyg-team/pytorch_geometric/pull/5668), [#5669](https://github.com/pyg-team/pytorch_geometric/pull/5669), [#5673](https://github.com/pyg-team/pytorch_geometric/pull/5673), [#5675](https://github.com/pyg-team/pytorch_geometric/pull/5675), [#5673](https://github.com/pyg-team/pytorch_geometric/pull/5676), [#5678](https://github.com/pyg-team/pytorch_geometric/pull/5678), [#5682](https://github.com/pyg-team/pytorch_geometric/pull/5682), [#5683](https://github.com/pyg-team/pytorch_geometric/pull/5683), [#5684](https://github.com/pyg-team/pytorch_geometric/pull/5684), [#5685](https://github.com/pyg-team/pytorch_geometric/pull/5685), [#5687](https://github.com/pyg-team/pytorch_geometric/pull/5687), [#5688](https://github.com/pyg-team/pytorch_geometric/pull/5688), [#5695](https://github.com/pyg-team/pytorch_geometric/pull/5695), [#5699](https://github.com/pyg-team/pytorch_geometric/pull/5699), [#5701](https://github.com/pyg-team/pytorch_geometric/pull/5701), [#5702](https://github.com/pyg-team/pytorch_geometric/pull/5702), [#5703](https://github.com/pyg-team/pytorch_geometric/pull/5703), [#5706](https://github.com/pyg-team/pytorch_geometric/pull/5706), [#5707](https://github.com/pyg-team/pytorch_geometric/pull/5707), [#5710](https://github.com/pyg-team/pytorch_geometric/pull/5710), [#5714](https://github.com/pyg-team/pytorch_geometric/pull/5714), [#5715](https://github.com/pyg-team/pytorch_geometric/pull/5715), [#5722](https://github.com/pyg-team/pytorch_geometric/pull/5722), [#5725](https://github.com/pyg-team/pytorch_geometric/pull/5725), [#5726](https://github.com/pyg-team/pytorch_geometric/pull/5726), [#5729](https://github.com/pyg-team/pytorch_geometric/pull/5729), [#5730](https://github.com/pyg-team/pytorch_geometric/pull/5730), [#5732](https://github.com/pyg-team/pytorch_geometric/pull/5732), [#5733](https://github.com/pyg-team/pytorch_geometric/pull/5733), [#5743](https://github.com/pyg-team/pytorch_geometric/pull/5743), [#5734](https://github.com/pyg-team/pytorch_geometric/pull/5734), [#5735](https://github.com/pyg-team/pytorch_geometric/pull/5735),  [#5736](https://github.com/pyg-team/pytorch_geometric/pull/5736), [#5737](https://github.com/pyg-team/pytorch_geometric/pull/5737), [#5747](https://github.com/pyg-team/pytorch_geometric/pull/5747), [#5752](https://github.com/pyg-team/pytorch_geometric/pull/5752), [#5753](https://github.com/pyg-team/pytorch_geometric/pull/5753), [#5760](https://github.com/pyg-team/pytorch_geometric/pull/5760), [#5754](https://github.com/pyg-team/pytorch_geometric/pull/5754), [#5757](https://github.com/pyg-team/pytorch_geometric/pull/5757))
+- Improved type hint support ([#5603](https://github.com/pyg-team/pytorch_geometric/pull/5603), [#5659](https://github.com/pyg-team/pytorch_geometric/pull/5659), [#5664](https://github.com/pyg-team/pytorch_geometric/pull/5664), [#5665](https://github.com/pyg-team/pytorch_geometric/pull/5665), [#5666](https://github.com/pyg-team/pytorch_geometric/pull/5666), [#5667](https://github.com/pyg-team/pytorch_geometric/pull/5667), [#5668](https://github.com/pyg-team/pytorch_geometric/pull/5668), [#5669](https://github.com/pyg-team/pytorch_geometric/pull/5669), [#5673](https://github.com/pyg-team/pytorch_geometric/pull/5673), [#5675](https://github.com/pyg-team/pytorch_geometric/pull/5675), [#5673](https://github.com/pyg-team/pytorch_geometric/pull/5676), [#5678](https://github.com/pyg-team/pytorch_geometric/pull/5678), [#5682](https://github.com/pyg-team/pytorch_geometric/pull/5682), [#5683](https://github.com/pyg-team/pytorch_geometric/pull/5683), [#5684](https://github.com/pyg-team/pytorch_geometric/pull/5684), [#5685](https://github.com/pyg-team/pytorch_geometric/pull/5685), [#5687](https://github.com/pyg-team/pytorch_geometric/pull/5687), [#5688](https://github.com/pyg-team/pytorch_geometric/pull/5688), [#5695](https://github.com/pyg-team/pytorch_geometric/pull/5695), [#5699](https://github.com/pyg-team/pytorch_geometric/pull/5699), [#5701](https://github.com/pyg-team/pytorch_geometric/pull/5701), [#5702](https://github.com/pyg-team/pytorch_geometric/pull/5702), [#5703](https://github.com/pyg-team/pytorch_geometric/pull/5703), [#5706](https://github.com/pyg-team/pytorch_geometric/pull/5706), [#5707](https://github.com/pyg-team/pytorch_geometric/pull/5707), [#5710](https://github.com/pyg-team/pytorch_geometric/pull/5710), [#5714](https://github.com/pyg-team/pytorch_geometric/pull/5714), [#5715](https://github.com/pyg-team/pytorch_geometric/pull/5715), [#5722](https://github.com/pyg-team/pytorch_geometric/pull/5722), [#5724](https://github.com/pyg-team/pytorch_geometric/pull/5724), [#5725](https://github.com/pyg-team/pytorch_geometric/pull/5725), [#5726](https://github.com/pyg-team/pytorch_geometric/pull/5726), [#5729](https://github.com/pyg-team/pytorch_geometric/pull/5729), [#5730](https://github.com/pyg-team/pytorch_geometric/pull/5730), [#5732](https://github.com/pyg-team/pytorch_geometric/pull/5732), [#5733](https://github.com/pyg-team/pytorch_geometric/pull/5733), [#5743](https://github.com/pyg-team/pytorch_geometric/pull/5743), [#5734](https://github.com/pyg-team/pytorch_geometric/pull/5734), [#5735](https://github.com/pyg-team/pytorch_geometric/pull/5735), [#5736](https://github.com/pyg-team/pytorch_geometric/pull/5736), [#5737](https://github.com/pyg-team/pytorch_geometric/pull/5737), [#5747](https://github.com/pyg-team/pytorch_geometric/pull/5747), [#5752](https://github.com/pyg-team/pytorch_geometric/pull/5752), [#5753](https://github.com/pyg-team/pytorch_geometric/pull/5753), [#5760](https://github.com/pyg-team/pytorch_geometric/pull/5760), [#5754](https://github.com/pyg-team/pytorch_geometric/pull/5754), [#5757](https://github.com/pyg-team/pytorch_geometric/pull/5757))
 - 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_grid.py

@@ -1,3 +1,6 @@
+import torch
+
+from torch_geometric.testing import is_full_test
 from torch_geometric.utils import grid
 
 
@@ -14,3 +17,10 @@ def test_grid():
     assert row.tolist() == expected_row
     assert col.tolist() == expected_col
     assert pos.tolist() == expected_pos
+
+    if is_full_test():
+        jit = torch.jit.script(grid)
+        (row, col), pos = jit(height=3, width=2)
+        assert row.tolist() == expected_row
+        assert col.tolist() == expected_col
+        assert pos.tolist() == expected_pos

torch_geometric/utils/coalesce.py

@@ -7,9 +7,30 @@
 from .num_nodes import maybe_num_nodes
 
 
+@torch.jit._overload
+def coalesce(edge_index, edge_attr=None, num_nodes=None, reduce="add",
+             is_sorted=False, sort_by_row=True):
+    # type: (Tensor, Optional[bool], Optional[int], str, bool, bool) -> Tensor  # noqa
+    pass
+
+
+@torch.jit._overload
+def coalesce(edge_index, edge_attr=None, num_nodes=None, reduce="add",
+             is_sorted=False, sort_by_row=True):
+    # type: (Tensor, Tensor, Optional[int], str, bool, bool) -> Tuple[Tensor, Tensor]  # noqa
+    pass
+
+
+@torch.jit._overload
+def coalesce(edge_index, edge_attr=None, num_nodes=None, reduce="add",
+             is_sorted=False, sort_by_row=True):
+    # type: (Tensor, List[Tensor], Optional[int], str, bool, bool) -> Tuple[Tensor, List[Tensor]]  # noqa
+    pass
+
+
 def coalesce(
     edge_index: Tensor,
-    edge_attr: Optional[Union[Tensor, List[Tensor]]] = None,
+    edge_attr: Union[Optional[Tensor], List[Tensor]] = None,
     num_nodes: Optional[int] = None,
     reduce: str = "add",
     is_sorted: bool = False,
@@ -74,18 +95,21 @@ def coalesce(
     if not is_sorted:
         idx[1:], perm = idx[1:].sort()
         edge_index = edge_index[:, perm]
-        if edge_attr is not None and isinstance(edge_attr, Tensor):
+        if isinstance(edge_attr, Tensor):
             edge_attr = edge_attr[perm]
-        elif edge_attr is not None:
+        elif isinstance(edge_attr, (list, tuple)):
             edge_attr = [e[perm] for e in edge_attr]
 
     mask = idx[1:] > idx[:-1]
 
     # Only perform expensive merging in case there exists duplicates:
     if mask.all():
-        return edge_index if edge_attr is None else (edge_index, edge_attr)
+        if isinstance(edge_attr, (Tensor, list, tuple)):
+            return edge_index, edge_attr
+        return edge_index
 
     edge_index = edge_index[:, mask]
+
     if edge_attr is None:
         return edge_index
 
@@ -95,9 +119,11 @@ def coalesce(
 
     if isinstance(edge_attr, Tensor):
         edge_attr = scatter(edge_attr, idx, 0, None, dim_size, reduce)
-    else:
+        return edge_index, edge_attr
+    elif isinstance(edge_attr, (list, tuple)):
         edge_attr = [
             scatter(e, idx, 0, None, dim_size, reduce) for e in edge_attr
         ]
+        return edge_index, edge_attr
 
-    return edge_index, edge_attr
+    return edge_index

torch_geometric/utils/grid.py

@@ -1,8 +1,17 @@
+from typing import Optional, Tuple
+
 import torch
-from torch_sparse import coalesce
+from torch import Tensor
+
+from torch_geometric.utils.coalesce import coalesce
 
 
-def grid(height, width, dtype=None, device=None):
+def grid(
+    height: int,
+    width: int,
+    dtype: Optional[torch.dtype] = None,
+    device: Optional[torch.device] = None,
+) -> Tuple[Tensor, Tensor]:
     r"""Returns the edge indices of a two-dimensional grid graph with height
     :attr:`height` and width :attr:`width` and its node positions.
 
@@ -29,13 +38,17 @@ def grid(height, width, dtype=None, device=None):
                 [0., 0.],
                 [1., 0.]])
     """
-
     edge_index = grid_index(height, width, device)
     pos = grid_pos(height, width, dtype, device)
     return edge_index, pos
 
 
-def grid_index(height, width, device=None):
+def grid_index(
+    height: int,
+    width: int,
+    device: Optional[torch.device] = None,
+) -> Tensor:
+
     w = width
     kernel = [-w - 1, -1, w - 1, -w, 0, w, -w + 1, 1, w + 1]
     kernel = torch.tensor(kernel, device=device)
@@ -51,12 +64,17 @@ def grid_index(height, width, device=None):
     row, col = row[mask], col[mask]
 
     edge_index = torch.stack([row, col], dim=0)
-    edge_index, _ = coalesce(edge_index, None, height * width, height * width)
-
+    edge_index = coalesce(edge_index, None, height * width)
     return edge_index
 
 
-def grid_pos(height, width, dtype=None, device=None):
+def grid_pos(
+    height: int,
+    width: int,
+    dtype: Optional[torch.dtype] = None,
+    device: Optional[torch.device] = None,
+) -> Tensor:
+
     dtype = torch.float if dtype is None else dtype
     x = torch.arange(width, dtype=dtype, device=device)
     y = (height - 1) - torch.arange(height, dtype=dtype, device=device)