[Code Coverage] loader/utils.py (#7857)

Part of #6528.

---------

Co-authored-by: rusty1s <matthias.fey@tu-dortmund.de>

test/loader/test_utils.py

@@ -0,0 +1,16 @@
+import pytest
+import torch
+
+from torch_geometric.loader.utils import index_select
+
+
+def test_index_select():
+    x = torch.randn(3, 5)
+    index = torch.tensor([0, 2])
+    assert torch.equal(index_select(x, index), x[index])
+    assert torch.equal(index_select(x, index, dim=-1), x[..., index])
+
+
+def test_index_select_out_of_range():
+    with pytest.raises(IndexError, match="out of range"):
+        index_select(torch.randn(3, 5), torch.tensor([0, 2, 3]))

torch_geometric/loader/utils.py

@@ -29,10 +29,29 @@
 )
 
 
-def index_select(value: FeatureTensorType, index: Tensor,
-                 dim: int = 0) -> Tensor:
-
-    # PyTorch currently only supports indexing via `torch.int64` :(
+def index_select(
+    value: FeatureTensorType,
+    index: Tensor,
+    dim: int = 0,
+) -> Tensor:
+    r"""Indexes the :obj:`value` tensor along dimension :obj:`dim` using the
+    entries in :obj:`index`.
+
+    Args:
+        value (torch.Tensor or np.ndarray): The input tensor.
+        index (torch.Tensor): The 1-D tensor containing the indices to index.
+        dim (int, optional): The dimension in which to index.
+            (default: :obj:`0`)
+
+    .. warning::
+
+        :obj:`index` is casted to a :obj:`torch.int64` tensor internally, as
+        `PyTorch currently only supports indexing
+        <https://github.com/pytorch/pytorch/issues/61819>`_ via
+        :obj:`torch.int64`.
+    """
+    # PyTorch currently only supports indexing via `torch.int64`:
+    # https://github.com/pytorch/pytorch/issues/61819
     index = index.to(torch.int64)
 
     if isinstance(value, Tensor):