Repository: dmlc/dgl
Pull request: [Dataloading] Ignore edge types without reverse types in edge dataloader
## Description
Fixes https://github.com/dmlc/dgl/issues/5410


## Checklist
Please feel free to remove inapplicable items for your PR.
- [x] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
- [x] I've leverage the [tools](https://docs.google.com/document/d/1iHyj7zlmygKSk5gBPsqIqL5ASPzJSPREaNT_QdsiYA4/edit) to beautify the python and c++ code.
- [x] The PR is complete and small, read the [Google eng practice (CL equals to PR)](https://google.github.io/eng-practices/review/developer/small-cls.html) 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).
- [x] All changes have test coverage
- [x] Code is well-documented
- [x] To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
- [x] Related issue is referred in this PR




Problem statement:
[DGL] Edge sampler failures with mixed reverses.
## 🐛 Bug

Edge prediction sampler fails when some edges have reverses and others do not.

## To Reproduce

```
import dgl
import torch
data_dict = {('A','AB','B'): (torch.tensor([0,1]),torch.tensor([0,1])),
             ('B','BA','A'): (torch.tensor([0,1]),torch.tensor([0,1])),
            ('B','BC','C'):(torch.tensor([0]),torch.tensor([0])),
            ('C','CA','A'): (torch.tensor([0,1]),torch.tensor([0,1]))}
g = dgl.heterograph(data_dict=data_dict)
block_sampler = dgl.dataloading.MultiLayerNeighborSampler(
            fanouts=[1],
            replace=True)
block_sampler = dgl.dataloading.as_edge_prediction_sampler(
            block_sampler,
            exclude="reverse_types",
            reverse_etypes={'AB':'BA'},

        )
d = dgl.dataloading.DataLoader(
            graph=g,
            indices={'AB':torch.tensor([0]),'BC':torch.tensor([0]),},
            graph_sampler=block_sampler,
            batch_size=2,
            shuffle=True,
            drop_last=False,
            num_workers=0,
            device='cuda',
            use_ddp=False,
            use_uva=True,
        )

next(iter(d))
```


## Expected behavior

The dataloader not to fail.

## Environment

Reproducible in multiple environments

## Additional context

Hints from issue discussion:
This could be solved by performing a get check on the dictionary being accessed: 

https://github.com/dmlc/dgl/blob/325e795a18bec66595bca59708b42b726c8f8d22/python/dgl/dataloading/base.py#L276

Commit message: Merge branch 'master' into fix-reverse-etypes

Patch to review:
diff --git a/python/dgl/dataloading/base.py b/python/dgl/dataloading/base.py
index d34702290511..5fcb70b87ff7 100644
--- a/python/dgl/dataloading/base.py
+++ b/python/dgl/dataloading/base.py
@@ -280,7 +280,11 @@ def _find_exclude_eids_with_reverse_types(g, eids, reverse_etype_map):
         for k, v in reverse_etype_map.items()
     }
     exclude_eids.update(
-        {reverse_etype_map[k]: v for k, v in exclude_eids.items()}
+        {
+            reverse_etype_map[k]: v
+            for k, v in exclude_eids.items()
+            if k in reverse_etype_map
+        }
     )
     return exclude_eids
 
diff --git a/tests/python/pytorch/dataloading/test_dataloader.py b/tests/python/pytorch/dataloading/test_dataloader.py
index 89ec63f0bb7e..e48ec00a6f00 100644
--- a/tests/python/pytorch/dataloading/test_dataloader.py
+++ b/tests/python/pytorch/dataloading/test_dataloader.py
@@ -622,6 +622,36 @@ def test_edge_dataloader_excludes(
             break
 
 
+def test_edge_dataloader_exclusion_without_all_reverses():
+    data_dict = {('A','AB','B'): (torch.tensor([0,1]),torch.tensor([0,1])),
+             ('B','BA','A'): (torch.tensor([0,1]),torch.tensor([0,1])),
+            ('B','BC','C'):(torch.tensor([0]),torch.tensor([0])),
+            ('C','CA','A'): (torch.tensor([0,1]),torch.tensor([0,1]))}
+    g = dgl.heterograph(data_dict=data_dict)
+    block_sampler = dgl.dataloading.MultiLayerNeighborSampler(
+                fanouts=[1],
+                replace=True)
+    block_sampler = dgl.dataloading.as_edge_prediction_sampler(
+                block_sampler,
+                exclude="reverse_types",
+                reverse_etypes={'AB':'BA'},
+
+            )
+    d = dgl.dataloading.DataLoader(
+                graph=g,
+                indices={'AB':torch.tensor([0]),'BC':torch.tensor([0]),},
+                graph_sampler=block_sampler,
+                batch_size=2,
+                shuffle=True,
+                drop_last=False,
+                num_workers=0,
+                device='cuda',
+                use_ddp=False,
+                use_uva=True,
+            )
+
+    next(iter(d))
+    
 def dummy_worker_init_fn(worker_id):
     pass
 
@@ -641,9 +671,9 @@ def test_dataloader_worker_init_fn():
     for _ in dataloader:
         pass
 
-
 if __name__ == "__main__":
     # test_node_dataloader(F.int32, 'neighbor', None)
     test_edge_dataloader_excludes(
         "reverse_types", False, 1, dgl.dataloading.ShaDowKHopSampler([5])
     )
+    test_edge_dataloader_exclusion_without_all_reverses()


Code review comments:
diff_hunk: @@ -622,6 +622,36 @@ def test_edge_dataloader_excludes(
             break
 
 
+def test_edge_dataloader_exclusion_without_all_reverses():
+    data_dict = {('A','AB','B'): (torch.tensor([0,1]),torch.tensor([0,1])),
+             ('B','BA','A'): (torch.tensor([0,1]),torch.tensor([0,1])),
+            ('B','BC','C'):(torch.tensor([0]),torch.tensor([0])),
+            ('C','CA','A'): (torch.tensor([0,1]),torch.tensor([0,1]))}
+    g = dgl.heterograph(data_dict=data_dict)
+    block_sampler = dgl.dataloading.MultiLayerNeighborSampler(
+                fanouts=[1],
+                replace=True)
+    block_sampler = dgl.dataloading.as_edge_prediction_sampler(
+                block_sampler,
+                exclude="reverse_types",
+                reverse_etypes={'AB':'BA'},
+
+            )
+    d = dgl.dataloading.DataLoader(
+                graph=g,
+                indices={'AB':torch.tensor([0]),'BC':torch.tensor([0]),},
+                graph_sampler=block_sampler,
+                batch_size=2,
+                shuffle=True,
+                drop_last=False,
+                num_workers=0,
+                device='cuda',
path: tests/python/pytorch/dataloading/test_dataloader.py
text: @user1:
The test is for CUDA only. My understanding is that this change applies to both CPU and GPU, so consider set the device according to the global context.