[Community Sprint] Add missing type hints and TorchScript support 🚀

[rusty1s]

🚀 The feature, motivation and pitch

We are kicking off our very first community sprint!

The community sprint resolves around adding missing type hints and TorchScript support for various functions across PyG, aiming to improve and clean-up our core codebase. Each individual contribution is designed to only take around 30 minutes to two hours to complete.

The sprint begins Wednesday October 12th with a kick off meeting at 8am PST. The community sprint will last 2 weeks and we will have another live hangouts when the sprint has completed. If you are interested in helping out, please also join our PyG slack channel #community-sprint-type-hints for more information. You can assign yourself to the type hints you are planning to work on here.

🚀 Add missing type hints and TorchScript support

Type hints are currently used inconsistently in the torch-geometric repository, and it would be nice to make them a complete, consistent thing across all datasets, models and utilities. Adding type hint support in models also helps us to improve our TorchScript coverage across layers and models provided in nn.*.

Example

Take a look at the current implementation of contains_isolated_nodes:

def contains_isolated_nodes(edge_index, num_nodes=None):
    num_nodes = maybe_num_nodes(edge_index, num_nodes)
    edge_index, _ = remove_self_loops(edge_index)
    return torch.unique(edge_index.view(-1)).numel() < num_nodes

Adding type hints support to the function signature helps us to better understand its input and output, improving code readability:

def contains_isolated_nodes(
   edge_index: Tensor,
   num_nodes: Optional[int] = None,
) -> bool:
   ...

Importantly, it also lets us use it as part of a TorchScript model. Without it, all arguments that miss type hints are expected to be PyTorch tensors (which is clearly not the case for the num_nodes argument). Without it, torch.jit.script compilation will fail:

import torch

from torch_geometric.utils import contains_isolated_nodes

contains_isolated_nodes = torch.jit.script(contains_isolated_nodes)

contains_isolated_nodes(torch.tensor([[0, 1, 0], [1, 0, 0]])

RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File ".../pytorch_geometric/torch_geometric/utils/isolated.py", line 29, in contains_isolated_nodes
    num_nodes = maybe_num_nodes(edge_index, num_nodes)
                ~~~~~~~~~~~~~~~ <--- HERE
    edge_index, _ = remove_self_loops(edge_index)
    return torch.unique(edge_index.view(-1)).numel() < num_nodes
RuntimeError: Cannot input a tensor of dimension other than 0 as a scalar argument

Guide to contributing

See here for a basic example to follow.

1. Ensure you have read our contributing guidelines.
2. Claim the functionality/model you want to improve here.
3. Implement the changes as in [Type Hints] utils.contains_isolated_nodes #5603. At best, ensure in a test that the model/function is convertable to a TorchScript program and that it results in the same output. It is okay to add type hint support for multiple functions/models within a single PR as long as you have assigned yourself to each of them, and the number of file changes stays at a reasonable number to ease reviewing (e.g., not more than 10 touched files).
4. Open a PR to the PyG repository and name it: "[Type Hints] {model_name/function_name}". In addition, ensure that documentation is rendered properly (CI will build the documentation automatically for you). Afterwards, add your PR number to the "Improved type hint support" line in CHANGELOG.md.

Tips for making your PR

• If you are unfamiliar with how type hints work, you can read the Python library documentation on them, but it is probably even easier to just look at another PR that added them.
• The types will usually be obvious, e.g., Tensor, bool, int, float. Wrap them within Optional[*] whenever the argument can be None.
• Specialized PyG type hints (e.g., Adj) are defined in typing.py.
• In some rare cases, type hints are challenging to add, e.g., whenever a model/function supports a Union of different types (which may be the case for functions/models that also support SparseTensor, e.g., edge_index: Union[Tensor, SparseTensor]). In that case, TorchScript support can be achieved via the torch.jit._overload decorator. See here for an example.
• The corresponding tests of PyG models and functions can be found in the test/ directory. For example, tests for torch_geometric/utils/isolated.py can be found in test/utils/test_isolated.py. You can run individual test files via pytest test/utils/test_isolated.py. It is only necessary to test TorchScript support for the nn.* and utils.* packages. No changes necessary for datasets.* and transforms.* packages.
• Ensure that the TorchScript variant compiles and achieves the same output:

  from torch_geometric.testing import is_full_test
  
  ...
  
  edge_index = torch.tensor([[0, 1, 2, 0], [1, 0, 2, 0]])
  assert contains_isolated_nodes(edge_index)
    
  if is_full_test():
      jit = torch.jit.script(contains_isolated_nodes)
      assert jit(edge_index)

  Note that we generally gate TorchScript tests behind is_full_test() which guarantees that TorchScript tests are only run nightly. You can enable full tests locally via the FULL_TEST=1 environment variable, e.g., FULL_TEST=1 pytest test/utils/test_isolated.py. It is only necessary to ensure TorchScript support for the nn.* and utils.* packages. No changes necessary for datasets.* and transforms.* packages.

Functions/Models to update

This list may be incomplete. If you still find a function without missing type hints/TorchScript tests, please let us know or add them on your own.

• nn.MetaLayer ([Type Hints] nn.MetaLayer #5758)
• nn.ChebConv ([Type Hints] nn.ChebConv #5730)
• nn.InstanceNorm ([Type Hints] nn.InstanceNorm #5684)
• nn.GraphSizeNorm ([Type Hints] nn.GraphSizeNorm #5729)
• nn.MessageNorm ([Type Hints] nn.MessageNorm #5847)
• nn.DiffGroupNorm ([Type Hints] nn.DiffGroupNorm #5757)
• nn.TopKPooling ([Type Hints] nn.TopKPooling #5731)
• nn.SAGPooling ([Type Hints] nn.SAGPooling #5810)
• nn.EdgePooling ([Type Hints] nn.EdgePooling #5738)
• nn.PANPooling ([Type Hints] nn.PANPooling #5852)
• nn.max_pool ([Type Hints] nn.max_pool #5735)
• nn.avg_pool ([Type Hints] nn.avg_pool #5734)
• nn.max_pool_x ([Type Hints] nn.max_pool_x #5687)
• nn.max_pool_neighbor_x ([Type Hints] nn.max_pool_neighbor_x #5703)
• nn.avg_pool_x ([Type Hints] nn.avg_pool_x #5706)
• nn.avg_pool_neighbor_x ([Type Hints] nn.avg_pool_neighbor_x #5707)
• nn.Node2Vec ([TypeHints] Node2Vec #5669)
• nn.DeepGraphInfomax ([Type Hints] nn.DeepGraphInfomax #5688)
• nn.InnerProductDecoder ([Type Hints] nn.DenseGraphConv, DenseGINConv, GAE, VGAE #5842)
• nn.GAE ([Type Hints] nn.DenseGraphConv, DenseGINConv, GAE, VGAE #5842)
• nn.VGAE ([Type Hints] nn.DenseGraphConv, DenseGINConv, GAE, VGAE #5842)
• nn.ARGA ([Type hints] nn.ARGVA and nn.ARGA #5726)
• nn.ARGVA ([Type hints] nn.ARGVA and nn.ARGA #5726)
• nn.SignedGCN ([Type hints] nn.SignedGCN #5725)
• nn.RENet ([Type Hints] nn.RENet #5715)
• nn.GraphUNet ([Type Hints] nn.GraphUNet and nn.SchNet #5710)
• nn.SchNet ([Type Hints] nn.GraphUNet and nn.SchNet #5710)
• nn.DimeNet ([Type Hints] nn.DimeNet #5806)
• nn.DimeNetPlusPlus ([Type Hints] nn.DimeNet #5806)
• nn.GNNExplainer ([Type Hints] nn.GNNExplainer #5716)
• nn.DeepGCNLayer ([Type Hints] nn.DeepGCNLayer #5699)
• nn.AttentiveFP([Type Hints] nn.AttentiveFP #5766)
• nn.DenseGCNConv ([Type Hints] DenseSAGEConv and DenseGCNConv #5664)
• nn.DenseGINConv ([Type Hints] nn.DenseGraphConv, DenseGINConv, GAE, VGAE #5842)
• nn.DenseGraphConv ([Type Hints] nn.DenseGraphConv, DenseGINConv, GAE, VGAE #5842)
• nn.DenseSAGEConv ([Type Hints] DenseSAGEConv and DenseGCNConv #5664)
• nn.dense_diff_pool ([Type Hints] nn.dense_diff_pool #5754)
• nn.dense_mincut_pool ([Type Hints] nn.dense_mincut_pool #5756)
• datasets.NELL ([Type Hints] NELL PPI MoleculeNet QM7 ZINC #5678)
• datasets.PPI ([Type Hints] NELL PPI MoleculeNet QM7 ZINC #5678)
• datasets.Reddit ([Type Hints] datasets.Reddit and datasets.Reddit2 #5695)
• datasets.Reddit2 ([Type Hints] datasets.Reddit and datasets.Reddit2 #5695)
• datasets.Yelp ([Type Hints] datasets.YELP #5747)
• datasets.QM7b ([Type Hints] NELL PPI MoleculeNet QM7 ZINC #5678)
• datasets.ZINC ([Type Hints] NELL PPI MoleculeNet QM7 ZINC #5678)
• datasets.MoleculeNet ([Type Hints] NELL PPI MoleculeNet QM7 ZINC #5678)
• datasets.MNISTSuperpixels ([Type Hints] datasets.MNISTSuperpixels #5760)
• datasets.ShapeNet (Adding type hints for shapenet #5828)
• datasets.ModelNet ([Type Hints] datasets.ModelNet #5701)
• datasets.SHREC2016 ([Type Hints] datasets.SHREC2016 #5798)
• datasets.TOSCA ([Type Hints] datasets.TOSCA, MixHopSyntheticDataset, PCPNetDataset and JODIE #5797)
• datasets.PCPNetDataset ([Type Hints] datasets.TOSCA, MixHopSyntheticDataset, PCPNetDataset and JODIE #5797)
• datasets.S3DIS ([Type Hints] datasets.S3DIS #5799)
• datasets.ICEWS18 ([Type Hints] datasets.ICEWS18 #5666)
• datasets.WILLOWObjectClass ([Type Hints] datasets.WILLOWObjectClass, PascalVOCKeypoints and PascalPF #5781)
• datasets.PascalVOCKeypoints ([Type Hints] datasets.WILLOWObjectClass, PascalVOCKeypoints and PascalPF #5781)
• datasets.PascalPF ([Type Hints] datasets.WILLOWObjectClass, PascalVOCKeypoints and PascalPF #5781)
• datasets.SNAPDataset ([Type Hints] datasets.SNAPDataset, SuiteSparseMatrixCollection and WordNet18 #5811)
• datasets.SuiteSparseMatrixCollection ([Type Hints] datasets.SNAPDataset, SuiteSparseMatrixCollection and WordNet18 #5811)
• datasets.WordNet18 ([Type Hints] datasets.SNAPDataset, SuiteSparseMatrixCollection and WordNet18 #5811)
• datasets.WordNet18RR ([Type Hints] datasets.SNAPDataset, SuiteSparseMatrixCollection and WordNet18 #5811)
• datasets.WebKB ([Type Hints] datasets.WebKB #5778)
• datasets.JODIEDataset ([Type Hints] datasets.TOSCA, MixHopSyntheticDataset, PCPNetDataset and JODIE #5797)
• datasets.MixHopSyntheticDataset ([Type Hints] datasets.TOSCA, MixHopSyntheticDataset, PCPNetDataset and JODIE #5797)
• datasets.UPFD ([Type Hints] datasets.UPFD #5800)
• transforms.Distance ([Type Hints] transforms.Distance #5685)
• transforms.Cartesian ([Type Hints] transforms.Cartesian #5673)
• transforms.LocalCartesian ([Type Hints] transforms.LocalCartesian #5675)
• transforms.Polar ([Type Hints] transforms.Polar #5676)
• transforms.Spherical ([Type Hints] transforms.Spherical and transforms.PointPairFeatures #5732)
• transforms.PointPairFeatures ([Type Hints] transforms.Spherical and transforms.PointPairFeatures #5732)
• transforms.OneHotDegree ([TypeHints] OneHotDegree #5667)
• transforms.TargetIndegree ([Type Hints] transforms: TargetIndegree, ToSLIC, and SVDFeatureReduction #5743)
• transforms.RandomJitter ([TypeHints] Random Transforms #5714)
• transforms.RandomFlip ([TypeHints] Random Transforms #5714)
• transforms.RandomScale ([TypeHints] Random Transforms #5714)
• transforms.RandomRotate ([TypeHints] Random Transforms #5714)
• transforms.RandomShear ([Type Hints] transforms.RandomShear #5702)
• transforms.KNNGraph ([Type Hints] transforms.KNNGraph and transforms.FaceToEdge #5722)
• transforms.FaceToEdge ([Type Hints] transforms.KNNGraph and transforms.FaceToEdge #5722)
• transforms.SamplePoints ([Type Hints] transforms.SamplePoints, transforms.FixedPoints and transforms.LaplacianLambdaMax #5733)
• transforms.FixedPoints ([Type Hints] transforms.SamplePoints, transforms.FixedPoints and transforms.LaplacianLambdaMax #5733)
• transforms.ToDense ([TypeHints] ToDense #5668)
• transforms.LaplacianLambdaMax ([Type Hints] transforms.SamplePoints, transforms.FixedPoints and transforms.LaplacianLambdaMax #5733)
• transforms.ToSLIC ([Type Hints] transforms: TargetIndegree, ToSLIC, and SVDFeatureReduction #5743)
• transforms.GDC ([Type Hints] transforms.GDC #5752)
• transforms.SIGN ([Type Hints] transforms.SIGN, utils.to_scipy_sparse_matrix and utils.from_scipy_sparse_matrix #5736)
• transforms.SVDFeatureReduction ([Type Hints] transforms: TargetIndegree, ToSLIC, and SVDFeatureReduction #5743)
• transforms.AddSelfLoops ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.Center ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.Compose ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.Delaunay ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.GCNNorm ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.GenerateMeshNormals ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.LinearTransformation ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.LocalDegreeProfile ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.NormalizeFeatures ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.NormalizeRotation ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.NormalizeScale ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.RadiusGraph ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.RandomLinkSplit ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.RandomNodeSplit ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.RemoveIsolatedNodes ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.ToDevice ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.ToSparseTensor ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.ToUndirected ([Type Hints] Missing type hints in transforms.* #5753)
• transforms.TwoHop ([Type Hints] Missing type hints in transforms.* #5753)
• utils.remove_isolated_nodes ([Type Hints] utils.remove_isolated_nodes #5659)
• utils.get_laplacian ([Type Hints] utils.to_dense_adj and utils.get_laplacian #5682)
• utils.to_dense_adj ([Type Hints] utils.to_dense_adj and utils.get_laplacian #5682)
• utils.dense_to_sparse ([Type Hints] utils.dense_to_sparse #5683)
• utils.normalized_cut ([Type Hints] utils.normalized_cut #5665)
• utils.grid ([Type Hints] utils.grid #5724)
• utils.geodesic_distance ([Type Hints] utils.geodesic_distance and utils.tree_decomposition #5851)
• utils.tree_decomposition ([Type Hints] utils.geodesic_distance and utils.tree_decomposition #5851)
• utils.to_scipy_sparse_matrix ([Type Hints] transforms.SIGN, utils.to_scipy_sparse_matrix and utils.from_scipy_sparse_matrix #5736)
• utils.from_scipy_sparse_matrix ([Type Hints] transforms.SIGN, utils.to_scipy_sparse_matrix and utils.from_scipy_sparse_matrix #5736)
• utils.to_networkx ([Type Hints] to_networkx, from_networkx, erdos_renyi_graph, etc #5768)
• utils.from_networkx ([Type Hints] to_networkx, from_networkx, erdos_renyi_graph, etc #5768)
• utils.erdos_renyi_graph ([Type Hints] to_networkx, from_networkx, erdos_renyi_graph, etc #5768)
• utils.stochastic_blockmodel_graph ([Type Hints] to_networkx, from_networkx, erdos_renyi_graph, etc #5768)
• utils.barabasi_albert_graph ([Type Hints] to_networkx, from_networkx, erdos_renyi_graph, etc #5768)
• utils.train_test_split_edges ([Type Hints] utils.train_test_split_edges #5737)

Alternatives

No response

Additional context

No response

[EdisonLeeeee]

Looks great! I'm more than happy to take transforms.* :)

[rusty1s]

@EdisonLeeeee Sounds great. Please add yourself to the Google Sheet doc, but please leave room for other people to participate as well :)

[venomouscyanide]

@rusty1s is nn.InnerProductEncoder supposed to be nn.InnerProductDecoder ? Not sure where nn.InnerProductEncoder is present in the codebase.

[rusty1s]

@venomouscyanide This is indeed a typo, good catch. Fixed :)

[EdisonLeeeee]

@rusty1s Some of the classes in transforms.* is missing (partially) type hints, e.g., transforms.Center and transforms.AddSelfLoops, which are not mentioned in this issue. Shall we also add type hints for them?

[rusty1s]

Yeah, if you have found some transforms not included here, please add them to the list :)

[rusty1s]

We reached 100% coverage of the community sprint! 🎉 This is amazing to see! I like to deeply thank everyone who participated in this sprint. Thanks for improving our code base and for contributing to PyG!