[Type Hints] utils.grid

[nihal-rao]

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

[codecov]

Codecov Report

Merging #5724 (1439566) into master (7cb8c07) will decrease coverage by 0.00%.
The diff coverage is 84.00%.

❗ Current head 1439566 differs from pull request most recent head 9127a26. Consider uploading reports for the commit 9127a26 to get more accurate results

@@            Coverage Diff             @@
##           master    #5724      +/-   ##
==========================================
- Coverage   83.99%   83.98%   -0.01%     
==========================================
  Files         349      349              
  Lines       19353    19369      +16     
==========================================
+ Hits        16255    16267      +12     
- Misses       3098     3102       +4     

Impacted Files	Coverage Δ
torch_geometric/utils/coalesce.py	91.30% <77.77%> (-8.70%)	⬇️
torch_geometric/utils/grid.py	100.00% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[EdisonLeeeee]

Replacing

pytorch_geometric/torch_geometric/utils/coalesce.py

Lines 77 to 80 in 907d228

	if edge_attr is not None and isinstance(edge_attr, Tensor):
	edge_attr = edge_attr[perm]
	elif edge_attr is not None:
	edge_attr = [e[perm] for e in edge_attr]

with

if edge_attr is not None:
    if isinstance(edge_attr, Tensor):
        edge_attr = edge_attr[perm]
    else:
        edge_attr = [e[perm] for e in edge_attr]

should address the error.

[rusty1s]

I fixed this by adding TorchScript support to coalesce via overload decorators :)