MinCut pooling for sparse adjacency matrix

[xiaohan2012]

Motivation

The current nn.dense.mincut_pool requires the input adjacency matrix to be a dense matrix. This requirement does not scale to large graphs

What

This PR considers a sparse counterpart of nn.dense.mincut_pool, which takes in sparse adjacency matrix as input.

Preliminary experiments (not shown here) indicate that the sparse version runs much faster with smaller memory footprint.

Technical issues / questions

• unit tests for torch 1.13 failed with NotImplementedError: Could not run 'aten::sparse_dim' with arguments from the 'CPU' backend. I have little clue of how to fix it.

Remarks on 3D sparse-dense tensor multiplication

Currently, multiplying a batched 3D dense tensor (e.g., the node-level feature matrix) with a batched 3D sparse tensor (e.g., the adjacency matrix) is implemented by looping through the batches and performing 2D matrix multiplication (detailed in the _bmm function in this PR).

We can use a more efficient way to achieve the same goal:

1. diagonally stack the 2D batches in the sparse 3D tensor (e.g, using Batch.from_data_list),
2. vertically stack the 2D batches in the dense tensor, and
3. do the multiplication on the stacked 2D matrices (elaborated here)

Other remark

This PR is WIP. It is my first PR for this project, appreciate any feedback 🙏🙏🙏

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 78.83%. Comparing base (01e8303) to head (dc9cb52).
Report is 7 commits behind head on master.

❗ Current head dc9cb52 differs from pull request most recent head 156a03c. Consider uploading reports for the commit 156a03c to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #9243      +/-   ##
==========================================
- Coverage   85.22%   78.83%   -6.40%     
==========================================
  Files         459      471      +12     
  Lines       29951    30318     +367     
==========================================
- Hits        25526    23901    -1625     
- Misses       4425     6417    +1992     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.