Repository: dmlc/dgl

Issue:
Data type context check on torch.int32 graph raise DGLError
## 🐛 Bug

Hi there,

I recently found a bug when using dgl.dataloading.NodeDataLoader on graph constructed with torch.int32 node tensor.
By adding additional print() in dgl source files, eg: /dgl/sampling/neighbor.py line 311, it seems dgl created seed nodes tensor with dtype torch.int64, then compared it with my graph idtype, and raised context unequal error.

I could reproduce this bug on official dataset ACM.mat

## To Reproduce

Steps to reproduce the behavior:

Code sample:
```
import scipy.io
import urllib.request
import dgl
import torch

data_url = 'https://data.dgl.ai/dataset/ACM.mat'
data_file_path = '/tmp/ACM.mat'

urllib.request.urlretrieve(data_url, data_file_path)
data = scipy.io.loadmat(data_file_path)
print(list(data.keys()))

G = dgl.heterograph({
        ('paper', 'written-by', 'author') : data['PvsA'].nonzero(),
        ('author', 'writing', 'paper') : data['PvsA'].transpose().nonzero(),
        ('paper', 'citing', 'paper') : data['PvsP'].nonzero(),
        ('paper', 'cited', 'paper') : data['PvsP'].transpose().nonzero(),
        ('paper', 'is-about', 'subject') : data['PvsL'].nonzero(),
        ('subject', 'has', 'paper') : data['PvsL'].transpose().nonzero(),
    })

print(G)

G = G.int()
print(G.idtype) # got torch.int32 here

sampler = dgl.dataloading.MultiLayerNeighborSampler([10, 5])
dataloader = dgl.dataloading.NodeDataLoader(
                G,
                {ntype: torch.arange(G.number_of_nodes(ntype=ntype), dtype=torch.int32) for ntype in G.ntypes},
                sampler,
                batch_size=1024,
                shuffle=False,
                drop_last=False,
                num_workers=8
            )
for input_nodes, output_nodes, blocks in dataloader:
    print(input_nodes)   # boom
```
Stacktrace:
```
DGLError                                  Traceback (most recent call last)
Input In [7], in <cell line: 1>()
----> 1 for input_nodes, output_nodes, blocks in dataloader:
      2     print(input_nodes)

File ~/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/dataloading/dataloader.py:514, in _PrefetchingIter.__next__(self)
    512 def __next__(self):
    513     batch, feats, stream_event = \
--> 514         self._next_non_threaded() if not self.use_thread else self._next_threaded()
    515     batch = recursive_apply_pair(batch, feats, _assign_for)
    516     if stream_event is not None:

File ~/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/dataloading/dataloader.py:487, in _PrefetchingIter._next_non_threaded(self)
    486 def _next_non_threaded(self):
--> 487     batch = next(self.dataloader_it)
    488     batch = recursive_apply(batch, restore_parent_storage_columns, self.dataloader.graph)
    489     device = self.dataloader.device

File ~/anaconda3/envs/dgl8/lib/python3.8/site-packages/torch/utils/data/dataloader.py:517, in _BaseDataLoaderIter.__next__(self)
    515 if self._sampler_iter is None:
    516     self._reset()
--> 517 data = self._next_data()
    518 self._num_yielded += 1
    519 if self._dataset_kind == _DatasetKind.Iterable and \
    520         self._IterableDataset_len_called is not None and \
    521         self._num_yielded > self._IterableDataset_len_called:

File ~/anaconda3/envs/dgl8/lib/python3.8/site-packages/torch/utils/data/dataloader.py:1199, in _MultiProcessingDataLoaderIter._next_data(self)
   1197 else:
   1198     del self._task_info[idx]
-> 1199     return self._process_data(data)

File ~/anaconda3/envs/dgl8/lib/python3.8/site-packages/torch/utils/data/dataloader.py:1225, in _MultiProcessingDataLoaderIter._process_data(self, data)
   1223 self._try_put_index()
   1224 if isinstance(data, ExceptionWrapper):
-> 1225     data.reraise()
   1226 return data

File ~/anaconda3/envs/dgl8/lib/python3.8/site-packages/torch/_utils.py:429, in ExceptionWrapper.reraise(self)
    425 elif getattr(self.exc_type, "message", None):
    426     # Some exceptions have first argument as non-str but explicitly
    427     # have message field
    428     raise self.exc_type(message=msg)
--> 429 raise self.exc_type(msg)

DGLError: Caught DGLError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 202, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 35, in fetch
    return self.collate_fn(data)
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 531, in __call__
    batch = self.sample_func(self.g, items)
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/dataloading/base.py", line 252, in sample
    result = self.sample_blocks(g, seed_nodes, exclude_eids=exclude_eids)
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/dataloading/neighbor_sampler.py", line 110, in sample_blocks
    frontier = g.sample_neighbors(
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/utils/internal.py", line 909, in _fn
    return func(*args, **kwargs)
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/sampling/neighbor.py", line 291, in sample_neighbors
    frontier = _sample_neighbors(
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/sampling/neighbor.py", line 311, in _sample_neighbors
    nodes = utils.prepare_tensor_dict(g, nodes, 'nodes')
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/utils/checks.py", line 71, in prepare_tensor_dict
    return {key : prepare_tensor(g, val, '{}["{}"]'.format(name, key))
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/utils/checks.py", line 71, in <dictcomp>
    return {key : prepare_tensor(g, val, '{}["{}"]'.format(name, key))
  File "/home/duyilin/anaconda3/envs/dgl8/lib/python3.8/site-packages/dgl/utils/checks.py", line 34, in prepare_tensor
    raise DGLError('Expect argument "{}" to have data type {} and device '
dgl._ffi.base.DGLError: Expect argument "nodes["author"]" to have data type torch.int32 and device context cpu. But got torch.int64 and cpu.
```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

Blocks get printed normally.

## Environment

 - DGL Version (e.g., 1.0): dgl-cu110 0.8.0.post1
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 1.8.1
 - OS (e.g., Linux): Linux
 - How you installed DGL (`conda`, `pip`, source): pip
 - Build command you used (if compiling from source): NA
 - Python version: 3.8
 - CUDA/cuDNN version (if applicable): 11.0
 - GPU models and configuration (e.g. V100): V100 32G
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:

Resolving patch:
diff --git a/python/dgl/dataloading/dataloader.py b/python/dgl/dataloading/dataloader.py
--- a/python/dgl/dataloading/dataloader.py
+++ b/python/dgl/dataloading/dataloader.py
@@ -21,7 +21,7 @@
 from .. import ndarray as nd
 from ..utils import (
     recursive_apply, ExceptionWrapper, recursive_apply_pair, set_num_threads,
-    create_shared_mem_array, get_shared_mem_array, context_of)
+    create_shared_mem_array, get_shared_mem_array, context_of, dtype_of)
 from ..frame import LazyFeature
 from ..storages import wrap_storage
 from .base import BlockSampler, as_edge_prediction_sampler
@@ -86,9 +86,11 @@ def __next__(self):
 
 
 def _get_id_tensor_from_mapping(indices, device, keys):
-    lengths = torch.LongTensor([
-        (indices[k].shape[0] if k in indices else 0) for k in keys]).to(device)
-    type_ids = torch.arange(len(keys), device=device).repeat_interleave(lengths)
+    dtype = dtype_of(indices)
+    lengths = torch.tensor(
+        [(indices[k].shape[0] if k in indices else 0) for k in keys],
+        dtype=dtype, device=device)
+    type_ids = torch.arange(len(keys), dtype=dtype, device=device).repeat_interleave(lengths)
     all_indices = torch.cat([indices[k] for k in keys if k in indices])
     return torch.stack([type_ids, all_indices], 1)
 
diff --git a/python/dgl/utils/internal.py b/python/dgl/utils/internal.py
--- a/python/dgl/utils/internal.py
+++ b/python/dgl/utils/internal.py
@@ -1019,4 +1019,8 @@ def context_of(data):
     else:
         return F.context(data)
 
+def dtype_of(data):
+    """Return the dtype of the data which can be either a tensor or a dict of tensors."""
+    return F.dtype(next(iter(data.values())) if isinstance(data, Mapping) else data)
+
 _init_api("dgl.utils.internal")


Test patch:
diff --git a/tests/pytorch/test_dataloader.py b/tests/pytorch/test_dataloader.py
--- a/tests/pytorch/test_dataloader.py
+++ b/tests/pytorch/test_dataloader.py
@@ -10,6 +10,7 @@
 from collections import defaultdict
 from collections.abc import Iterator, Mapping
 from itertools import product
+from test_utils import parametrize_dtype
 import pytest
 
 
@@ -89,6 +90,15 @@ def test_neighbor_nonuniform(num_workers):
         elif seed == 0:
             assert neighbors == {3, 4}
 
+def _check_dtype(data, dtype, attr_name):
+    if isinstance(data, dict):
+        for k, v in data.items():
+            assert getattr(v, attr_name) == dtype
+    elif isinstance(data, list):
+        for v in data:
+            assert getattr(v, attr_name) == dtype
+    else:
+        assert getattr(data, attr_name) == dtype
 
 def _check_device(data):
     if isinstance(data, dict):
@@ -100,10 +110,11 @@ def _check_device(data):
     else:
         assert data.device == F.ctx()
 
+@parametrize_dtype
 @pytest.mark.parametrize('sampler_name', ['full', 'neighbor', 'neighbor2'])
 @pytest.mark.parametrize('pin_graph', [False, True])
-def test_node_dataloader(sampler_name, pin_graph):
-    g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4]))
+def test_node_dataloader(idtype, sampler_name, pin_graph):
+    g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4])).astype(idtype)
     if F.ctx() != F.cpu() and pin_graph:
         g1.create_formats_()
         g1.pin_memory_()
@@ -123,13 +134,16 @@ def test_node_dataloader(sampler_name, pin_graph):
             _check_device(input_nodes)
             _check_device(output_nodes)
             _check_device(blocks)
+            _check_dtype(input_nodes, idtype, 'dtype')
+            _check_dtype(output_nodes, idtype, 'dtype')
+            _check_dtype(blocks, idtype, 'idtype')
 
     g2 = dgl.heterograph({
          ('user', 'follow', 'user'): ([0, 0, 0, 1, 1, 1, 2], [1, 2, 3, 0, 2, 3, 0]),
          ('user', 'followed-by', 'user'): ([1, 2, 3, 0, 2, 3, 0], [0, 0, 0, 1, 1, 1, 2]),
          ('user', 'play', 'game'): ([0, 1, 1, 3, 5], [0, 1, 2, 0, 2]),
          ('game', 'played-by', 'user'): ([0, 1, 2, 0, 2], [0, 1, 1, 3, 5])
-    })
+    }).astype(idtype)
     for ntype in g2.ntypes:
         g2.nodes[ntype].data['feat'] = F.copy_to(F.randn((g2.num_nodes(ntype), 8)), F.cpu())
     batch_size = max(g2.num_nodes(nty) for nty in g2.ntypes)
@@ -146,6 +160,9 @@ def test_node_dataloader(sampler_name, pin_graph):
         _check_device(input_nodes)
         _check_device(output_nodes)
         _check_device(blocks)
+        _check_dtype(input_nodes, idtype, 'dtype')
+        _check_dtype(output_nodes, idtype, 'dtype')
+        _check_dtype(blocks, idtype, 'idtype')
 
     if g1.is_pinned():
         g1.unpin_memory_()