[DistPart] Fix corner case in dist partition which always led to an assertion error being triggered.

[thvasilo]

Description

In some edge cases where workers end up with no rows of features to send over the network (e.g. in range partition) the existing code was creating tensors of shape (0, 0) and trying to communicate/aggregate those over the network, whereas workers that had non-zero rows assigned would communicate (num_rows, feature_dimension), leading to an assertion error being triggered in

dgl/tools/distpartitioning/gloo_wrapper.py

Lines 144 to 146 in 6f2ccbf

	assert (
	sizes[idx - 1][1:] == sizes[idx][1:]
	) # except first dimension remaining dimensions should all be the same

This PR handles the 2D tensor case as a special case, ensuring the correct shape for the tensors being sent over the network, (0, feature_dimension).

This PR also removes an array that was being created but never used, reducing the memory footprint of the function.

Ping @Rhett-Ying for review

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
• I've leverage the tools to beautify the python and c++ code.
• The PR is complete and small, read the Google eng practice (CL equals to PR) to understand more about small PR. In DGL, we consider PRs with less than 200 lines of core code change are small (example, test and documentation could be exempted).
• All changes have test coverage
• Code is well-documented
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
• Related issue is referred in this PR
• If the PR is for a new model/paper, I've updated the example index here.

Changes

[dgl-bot]

To trigger regression tests:

• @dgl-bot run [instance-type] [which tests] [compare-with-branch];
  For example: @dgl-bot run g4dn.4xlarge all dmlc/master or @dgl-bot run c5.9xlarge kernel,api dmlc/master

[dgl-bot]

Commit ID: 8971a87fc323285a29ef38fa6b185519209a6a22

Build ID: 1

Status: ✅ CI test succeeded.

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 5fc034371a9263198b37a20578aea1bf0da6c6fd

Build ID: 2

Status: ✅ CI test succeeded.

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: b7fb8403de6d46185981396770f507d20cfc2c12

Build ID: 3

Status: ✅ CI test succeeded.

Report path: link

Full logs path: link

[thvasilo]

Hi @Rhett-Ying fixed the lint errors, is this fine to merge now?