[Feature] extend sort_csr/csc_by_tag to edge

[Rhett-Ying]

Description

this feature is needed by the fix for sample_etype_neighbors in DistDGL: #4150.

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
• Changes are complete (i.e. I finished coding on this PR)
• 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: 4dadbb8abe57c1f09866e42e38ab8c6a7ccf7480

Build ID: 1

Status: ❌ CI test failed in Stage [Tensorflow CPU Unit test].

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 2f30d80

Build ID: 2

Status: ❌ CI test failed in Stage [Tensorflow CPU Unit test].

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 42b8cf9

Build ID: 3

Status: ❌ CI test failed in Stage [Tensorflow CPU Unit test].

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 42b8cf9

Build ID: 4

Status: ❌ CI test failed in Stage [Tensorflow CPU Unit test].

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 9724d9ffa36c8c9c00c6e311061b27e05564d108

Build ID: 5

Status: ✅ CI test succeeded

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: f334977

Build ID: 6

Status: ❌ CI test failed in Stage [Torch CPU Unit test].

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: f334977

Build ID: 7

Status: ✅ CI test succeeded

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 3900700

Build ID: 8

Status: ✅ CI test succeeded

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 9bb2fe789a37e115f1cc2e49dae4dfe7df01ad18

Build ID: 9

Status: ✅ CI test succeeded

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: d57a3aa1ec7d061ac98f511f24bfe9fb64aa7e60

Build ID: 10

Status: ❌ CI test failed in Stage [Lint Check].

Report path: link

Full logs path: link

[dgl-bot]

Commit ID: 2413f86

Build ID: 11

Status: ✅ CI test succeeded

Report path: link

Full logs path: link