Repository: pyg-team/pytorch_geometric
Pull request #9782 (2024-11-13 07:11:00)

Drop `TensorAttr.fully_specify`
Removes `TensorAttr.fully_specify` which was originally added in #4534.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9b98a9bed..73a782026 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Added PyTorch 2.5 support ([#9779](https://github.com/pyg-team/pytorch_geometric/pull/9779))
+- Added support for fast `Delaunay()` triangulation via the `torch_delaunay` package ([#9748](https://github.com/pyg-team/pytorch_geometric/pull/9748))
+- Added PyTorch 2.5 support ([#9779](https://github.com/pyg-team/pytorch_geometric/pull/9779), [#9779](https://github.com/pyg-team/pytorch_geometric/pull/9780))
 - Support 3D tetrahedral mesh elements of shape `[4, num_faces]` in the `FaceToEdge` transformation ([#9776](https://github.com/pyg-team/pytorch_geometric/pull/9776))
 - Added the `use_pcst` option to `WebQSPDataset` ([#9722](https://github.com/pyg-team/pytorch_geometric/pull/9722))
 - Allowed users to pass `edge_weight` to `GraphUNet` models ([#9737](https://github.com/pyg-team/pytorch_geometric/pull/9737))
@@ -21,6 +22,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed the `k_hop_subgraph()` method for directed graphs ([#9756](https://github.com/pyg-team/pytorch_geometric/pull/9756))
 - Fixed `utils.group_cat` concatenating dimension ([#9766](https://github.com/pyg-team/pytorch_geometric/pull/9766))
 - Fixed `WebQSDataset.process` raising exceptions ([#9665](https://github.com/pyg-team/pytorch_geometric/pull/9665))
 
diff --git a/README.md b/README.md
index c37a073bf..f96d72083 100644
--- a/README.md
+++ b/README.md
@@ -371,12 +371,12 @@ We recommend to start with a minimal installation, and install additional depend
 
 For ease of installation of these extensions, we provide `pip` wheels for all major OS/PyTorch/CUDA combinations, see [here](https://data.pyg.org/whl).
 
-#### PyTorch 2.4
+#### PyTorch 2.5
 
-To install the binaries for PyTorch 2.4.0, simply run
+To install the binaries for PyTorch 2.5.0, simply run
 
 ```
-pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+${CUDA}.html
+pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.5.0+${CUDA}.html
 ```
 
 where `${CUDA}` should be replaced by either `cpu`, `cu118`, `cu121`, or `cu124` depending on your PyTorch installation.
@@ -387,23 +387,23 @@ where `${CUDA}` should be replaced by either `cpu`, `cu118`, `cu121`, or `cu124`
 | **Windows** | ✅    | ✅      | ✅      | ✅      |
 | **macOS**   | ✅    |         |         |         |
 
-#### PyTorch 2.3
+#### PyTorch 2.4
 
-To install the binaries for PyTorch 2.3.0, simply run
+To install the binaries for PyTorch 2.4.0, simply run
 
 ```
-pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.3.0+${CUDA}.html
+pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+${CUDA}.html
 ```
 
-where `${CUDA}` should be replaced by either `cpu`, `cu118`, or `cu121` depending on your PyTorch installation.
+where `${CUDA}` should be replaced by either `cpu`, `cu118`, `cu121`, or `cu124` depending on your PyTorch installation.
 
-|             | `cpu` | `cu118` | `cu121` |
-| ----------- | ----- | ------- | ------- |
-| **Linux**   | ✅    | ✅      | ✅      |
-| **Windows** | ✅    | ✅      | ✅      |
-| **macOS**   | ✅    |         |         |
+|             | `cpu` | `cu118` | `cu121` | `cu124` |
+| ----------- | ----- | ------- | ------- | ------- |
+| **Linux**   | ✅    | ✅      | ✅      | ✅      |
+| **Windows** | ✅    | ✅      | ✅      | ✅      |
+| **macOS**   | ✅    |         |         |         |
 
-**Note:** Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0, PyTorch 1.7.0/1.7.1, PyTorch 1.8.0/1.8.1, PyTorch 1.9.0, PyTorch 1.10.0/1.10.1/1.10.2, PyTorch 1.11.0, PyTorch 1.12.0/1.12.1, PyTorch 1.13.0/1.13.1, PyTorch 2.0.0/2.0.1, PyTorch 2.1.0/2.1.1/2.1.2, and PyTorch 2.2.0/2.2.1/2.2.2 (following the same procedure).
+**Note:** Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0, PyTorch 1.7.0/1.7.1, PyTorch 1.8.0/1.8.1, PyTorch 1.9.0, PyTorch 1.10.0/1.10.1/1.10.2, PyTorch 1.11.0, PyTorch 1.12.0/1.12.1, PyTorch 1.13.0/1.13.1, PyTorch 2.0.0/2.0.1, PyTorch 2.1.0/2.1.1/2.1.2, PyTorch 2.2.0/2.2.1/2.2.2, and PyTorch 2.3.0/2.3.1 (following the same procedure).
 **For older versions, you might need to explicitly specify the latest supported version number** or install via `pip install --no-index` in order to prevent a manual installation from source.
 You can look up the latest supported version number [here](https://data.pyg.org/whl).
 
diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu
index 0cf402e2d..6dafba18a 100644
--- a/docker/Dockerfile.xpu
+++ b/docker/Dockerfile.xpu
@@ -27,5 +27,5 @@ RUN echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc
 
 # Install PyG
 RUN pip install ninja wheel ogb && pip install git+https://github.com/pyg-team/pyg-lib.git && \
-    pip install torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+cpu.html && \
+    pip install torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.5.0+cpu.html && \
     pip install torch_geometric
diff --git a/docs/source/install/installation.rst b/docs/source/install/installation.rst
index f4d0eadbc..c3a8bd3d9 100644
--- a/docs/source/install/installation.rst
+++ b/docs/source/install/installation.rst
@@ -72,7 +72,7 @@ For ease of installation of these extensions, we provide :obj:`pip` wheels for t
    .. code-block:: none
 
       python -c "import torch; print(torch.__version__)"
-      >>> 2.4.0
+      >>> 2.5.0
 
 #. Find the CUDA version :pytorch:`PyTorch` was installed with:
 
@@ -89,6 +89,7 @@ For ease of installation of these extensions, we provide :obj:`pip` wheels for t
 
    where :obj:`${TORCH}` and :obj:`${CUDA}` should be replaced by the specific :pytorch:`PyTorch` and CUDA versions, respectively:
 
+   * :pytorch:`PyTorch` 2.4: :obj:`${TORCH}=2.5.0` and :obj:`${CUDA}=cpu|cu118|cu121|cu124`
    * :pytorch:`PyTorch` 2.4: :obj:`${TORCH}=2.4.0` and :obj:`${CUDA}=cpu|cu118|cu121|cu124`
    * :pytorch:`PyTorch` 2.3: :obj:`${TORCH}=2.3.0` and :obj:`${CUDA}=cpu|cu118|cu121`
    * :pytorch:`PyTorch` 2.2: :obj:`${TORCH}=2.2.0` and :obj:`${CUDA}=cpu|cu118|cu121`
@@ -96,19 +97,19 @@ For ease of installation of these extensions, we provide :obj:`pip` wheels for t
    * :pytorch:`PyTorch` 2.0: :obj:`${TORCH}=2.0.0` and :obj:`${CUDA}=cpu|cu117|cu118`
    * :pytorch:`PyTorch` 1.13: :obj:`${TORCH}=1.13.0` and :obj:`${CUDA}=cpu|cu116|cu117`
 
-   For example, for :pytorch:`PyTorch` 2.4.* and CUDA 12.4, type:
+   For example, for :pytorch:`PyTorch` 2.5.* and CUDA 12.4, type:
 
    .. code-block:: none
 
-      pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+cu124.html
+      pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.5.0+cu124.html
 
-   For example, for :pytorch:`PyTorch` 2.3.* and CUDA 11.8, type:
+   For example, for :pytorch:`PyTorch` 2.4.* and CUDA 11.8, type:
 
    .. code-block:: none
 
-      pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.3.0+cu118.html
+      pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+cu118.html
 
-**Note:** Binaries of older versions are also provided for :pytorch:`PyTorch` 1.4.0, 1.5.0, 1.6.0, 1.7.0/1.7.1, 1.8.0/1.8.1, 1.9.0, 1.10.0/1.10.1/1.10.2, 1.11.0, 1.12.0/1.12.1, 1.13.0/1.13.1, 2.0.0/2.0.1, 2.1.0/2.1.1/2.1.2, and 2.2.0/2.2.1/2.2.2 (following the same procedure).
+**Note:** Binaries of older versions are also provided for :pytorch:`PyTorch` 1.4.0, 1.5.0, 1.6.0, 1.7.0/1.7.1, 1.8.0/1.8.1, 1.9.0, 1.10.0/1.10.1/1.10.2, 1.11.0, 1.12.0/1.12.1, 1.13.0/1.13.1, 2.0.0/2.0.1, 2.1.0/2.1.1/2.1.2, 2.2.0/2.2.1/2.2.2, and 2.3.0/2.3.1 (following the same procedure).
 **For older versions, you need to explicitly specify the latest supported version number** or install via :obj:`pip install --no-index` in order to prevent a manual installation from source.
 You can look up the latest supported version number `here <https://data.pyg.org/whl>`__.
 
diff --git a/docs/source/install/quick-start.html b/docs/source/install/quick-start.html
index 5aebbc3c9..c29305cbe 100644
--- a/docs/source/install/quick-start.html
+++ b/docs/source/install/quick-start.html
@@ -75,9 +75,9 @@
 
 <script>
   var torchList = [
+    ['torch-2.5.0', 'PyTorch 2.5.*'],
     ['torch-2.4.0', 'PyTorch 2.4.*'],
     ['torch-2.3.0', 'PyTorch 2.3.*'],
-    ['torch-2.2.0', 'PyTorch 2.2.*'],
   ];
 
   var osList = [
@@ -113,7 +113,7 @@
       $("#command pre").text('# macOS binaries do not support CUDA');
     }
 
-    else if (torch != "torch-2.4.0" && cuda == "cu124") {
+    else if (torch == "torch-2.3.0" && cuda == "cu124") {
       $("#command pre").text('# PyTorch version does not support CUDA 12.4');
     }
 
diff --git a/torch_geometric/data/feature_store.py b/torch_geometric/data/feature_store.py
index dbd7954bc..6674bf42e 100644
--- a/torch_geometric/data/feature_store.py
+++ b/torch_geometric/data/feature_store.py
@@ -74,13 +74,6 @@ class TensorAttr(CastMixin):
         r"""Whether the :obj:`TensorAttr` has no unset fields."""
         return all([self.is_set(key) for key in self.__dataclass_fields__])
 
-    def fully_specify(self) -> 'TensorAttr':
-        r"""Sets all :obj:`UNSET` fields to :obj:`None`."""
-        for key in self.__dataclass_fields__:
-            if not self.is_set(key):
-                setattr(self, key, None)
-        return self
-
     def update(self, attr: 'TensorAttr') -> 'TensorAttr':
         r"""Updates an :class:`TensorAttr` with set attributes from another
         :class:`TensorAttr`.
@@ -230,10 +223,11 @@ class AttrView(CastMixin):
 
             store[group_name, attr_name]()
         """
-        # Set all UNSET values to None:
-        out = copy.copy(self)
-        out._attr.fully_specify()
-        return out._store.get_tensor(out._attr)
+        attr = copy.copy(self._attr)
+        for key in attr.__dataclass_fields__:  # Set all UNSET values to None.
+            if not attr.is_set(key):
+                setattr(attr, key, None)
+        return self._store.get_tensor(attr)
 
     def __copy__(self) -> 'AttrView':
         out = self.__class__.__new__(self.__class__)
@@ -479,9 +473,7 @@ class FeatureStore(ABC):
         # CastMixin will handle the case of key being a tuple or TensorAttr
         # object:
         key = self._tensor_attr_cls.cast(key)
-        # We need to fully-specify the key for __setitem__ as it does not make
-        # sense to work with a view here:
-        key.fully_specify()
+        assert key.is_fully_specified()
         self.put_tensor(value, key)
 
     def __getitem__(self, key: TensorAttr) -> Any:
@@ -503,13 +495,16 @@ class FeatureStore(ABC):
         # If the view is not fully-specified, return a :class:`AttrView`:
         return self.view(attr)
 
-    def __delitem__(self, key: TensorAttr):
+    def __delitem__(self, attr: TensorAttr):
         r"""Supports :obj:`del store[tensor_attr]`."""
         # CastMixin will handle the case of key being a tuple or TensorAttr
         # object:
-        key = self._tensor_attr_cls.cast(key)
-        key.fully_specify()
-        self.remove_tensor(key)
+        attr = self._tensor_attr_cls.cast(attr)
+        attr = copy.copy(attr)
+        for key in attr.__dataclass_fields__:  # Set all UNSET values to None.
+            if not attr.is_set(key):
+                setattr(attr, key, None)
+        self.remove_tensor(attr)
 
     def __iter__(self):
         raise NotImplementedError
diff --git a/torch_geometric/transforms/delaunay.py b/torch_geometric/transforms/delaunay.py
index 8ea90333a..5c2d469ad 100644
--- a/torch_geometric/transforms/delaunay.py
+++ b/torch_geometric/transforms/delaunay.py
@@ -1,3 +1,5 @@
+from typing import List
+
 import torch
 
 from torch_geometric.data import Data
@@ -5,30 +7,78 @@ from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
 
 
+class _QhullTransform(BaseTransform):
+    r"""Q-hull implementation of delaunay triangulation."""
+    def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        import scipy.spatial
+
+        pos = data.pos.cpu().numpy()
+        tri = scipy.spatial.Delaunay(pos, qhull_options='QJ')
+        face = torch.from_numpy(tri.simplices)
+
+        data.face = face.t().contiguous().to(data.pos.device, torch.long)
+        return data
+
+
+class _ShullTransform(BaseTransform):
+    r"""Sweep-hull implementation of delaunay triangulation."""
+    def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        from torch_delaunay.functional import shull2d
+
+        face = shull2d(data.pos.cpu())
+        data.face = face.t().contiguous().to(data.pos.device)
+        return data
+
+
+class _SequentialTransform(BaseTransform):
+    r"""Runs the first successful transformation.
+
+    All intermediate exceptions are suppressed except the last.
+    """
+    def __init__(self, transforms: List[BaseTransform]) -> None:
+        assert len(transforms) > 0
+        self.transforms = transforms
+
+    def forward(self, data: Data) -> Data:
+        for i, transform in enumerate(self.transforms):
+            try:
+                return transform.forward(data)
+            except ImportError as e:
+                if i == len(self.transforms) - 1:
+                    raise e
+        return data
+
+
 @functional_transform('delaunay')
 class Delaunay(BaseTransform):
     r"""Computes the delaunay triangulation of a set of points
     (functional name: :obj:`delaunay`).
+
+    .. hint::
+        Consider installing the
+        `torch_delaunay <https://github.com/ybubnov/torch_delaunay>`_ package
+        to speed up computation.
     """
-    def forward(self, data: Data) -> Data:
-        import scipy.spatial
+    def __init__(self) -> None:
+        self._transform = _SequentialTransform([
+            _ShullTransform(),
+            _QhullTransform(),
+        ])
 
+    def forward(self, data: Data) -> Data:
         assert data.pos is not None
+        device = data.pos.device
 
         if data.pos.size(0) < 2:
-            data.edge_index = torch.tensor([], dtype=torch.long,
-                                           device=data.pos.device).view(2, 0)
-        if data.pos.size(0) == 2:
-            data.edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long,
-                                           device=data.pos.device)
+            data.edge_index = torch.empty(2, 0, dtype=torch.long,
+                                          device=device)
+        elif data.pos.size(0) == 2:
+            data.edge_index = torch.tensor([[0, 1], [1, 0]], device=device)
         elif data.pos.size(0) == 3:
-            data.face = torch.tensor([[0], [1], [2]], dtype=torch.long,
-                                     device=data.pos.device)
-        if data.pos.size(0) > 3:
-            pos = data.pos.cpu().numpy()
-            tri = scipy.spatial.Delaunay(pos, qhull_options='QJ')
-            face = torch.from_numpy(tri.simplices)
-
-            data.face = face.t().contiguous().to(data.pos.device, torch.long)
+            data.face = torch.tensor([[0], [1], [2]], device=device)
+        else:
+            data = self._transform.forward(data)
 
         return data
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index cf4bf7bd8..e63b1849b 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -15,6 +15,7 @@ WITH_PT22 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 2
 WITH_PT23 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 3
 WITH_PT24 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 4
 WITH_PT25 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 5
+WITH_PT26 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 6
 WITH_PT111 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 11
 WITH_PT112 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 12
 WITH_PT113 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 13
diff --git a/torch_geometric/utils/_subgraph.py b/torch_geometric/utils/_subgraph.py
index 5c2dcd77c..7540cc9f2 100644
--- a/torch_geometric/utils/_subgraph.py
+++ b/torch_geometric/utils/_subgraph.py
@@ -346,10 +346,12 @@ def k_hop_subgraph(
 
     subsets = [node_idx]
 
+    preserved_edge_mask = torch.zeros_like(edge_mask)
     for _ in range(num_hops):
         node_mask.fill_(False)
         node_mask[subsets[-1]] = True
         torch.index_select(node_mask, 0, row, out=edge_mask)
+        preserved_edge_mask |= edge_mask
         subsets.append(col[edge_mask])
 
     subset, inv = torch.cat(subsets).unique(return_inverse=True)
@@ -360,6 +362,8 @@ def k_hop_subgraph(
 
     if not directed:
         edge_mask = node_mask[row] & node_mask[col]
+    else:
+        edge_mask = preserved_edge_mask
 
     edge_index = edge_index[:, edge_mask]
 


Test patch:
diff --git a/test/data/test_feature_store.py b/test/data/test_feature_store.py
index c9a850966..23d284249 100644
--- a/test/data/test_feature_store.py
+++ b/test/data/test_feature_store.py
@@ -78,7 +78,7 @@ def test_feature_store():
 
     # Partially-specified forms, when called, produce a Tensor output
     # from the `TensorAttr` that has been partially specified.
-    store[group_name] = tensor
+    store[group_name, None, None] = tensor
     assert isinstance(store[group_name], AttrView)
     assert torch.equal(store[group_name](), tensor)
 
diff --git a/test/transforms/test_delaunay.py b/test/transforms/test_delaunay.py
index fc027a2ea..4fac3917c 100644
--- a/test/transforms/test_delaunay.py
+++ b/test/transforms/test_delaunay.py
@@ -5,30 +5,60 @@ from torch_geometric.testing import withPackage
 from torch_geometric.transforms import Delaunay
 
 
-@withPackage('scipy')
-def test_delaunay():
-    assert str(Delaunay()) == 'Delaunay()'
+def assert_one_point(transform: Delaunay) -> None:
+    data = Data(pos=torch.rand(1, 2))
+    data = transform(data)
+    assert len(data) == 2
+    assert data.edge_index.tolist() == [[], []]
 
-    pos = torch.tensor([[-1, -1], [-1, 1], [1, 1], [1, -1]], dtype=torch.float)
-    data = Data(pos=pos)
-    data = Delaunay()(data)
+
+def assert_two_points(transform: Delaunay) -> None:
+    data = Data(pos=torch.rand(2, 2))
+    data = transform(data)
     assert len(data) == 2
-    assert data.face.tolist() == [[3, 1], [1, 3], [0, 2]]
+    assert data.edge_index.tolist() == [[0, 1], [1, 0]]
 
-    pos = torch.tensor([[-1, -1], [-1, 1], [1, 1]], dtype=torch.float)
-    data = Data(pos=pos)
-    data = Delaunay()(data)
+
+def assert_three_points(transform: Delaunay) -> None:
+    data = Data(pos=torch.rand(3, 2))
+    data = transform(data)
     assert len(data) == 2
     assert data.face.tolist() == [[0], [1], [2]]
 
-    pos = torch.tensor([[-1, -1], [1, 1]], dtype=torch.float)
-    data = Data(pos=pos)
-    data = Delaunay()(data)
-    assert len(data) == 2
-    assert data.edge_index.tolist() == [[0, 1], [1, 0]]
 
-    pos = torch.tensor([[-1, -1]], dtype=torch.float)
+def assert_four_points(transform: Delaunay) -> None:
+    pos = torch.tensor([
+        [-1.0, -1.0],
+        [-1.0, 1.0],
+        [1.0, 1.0],
+        [0.5, -0.5],
+    ])
     data = Data(pos=pos)
-    data = Delaunay()(data)
+    data = transform(data)
     assert len(data) == 2
-    assert data.edge_index.tolist() == [[], []]
+
+    # The order of the simplices does not matter, therefore assert sets.
+    faces = set(map(tuple, data.face.tolist()))
+    assert faces == {(3, 1), (1, 3), (0, 2)}
+
+
+@withPackage('scipy')
+def test_qhull_delaunay() -> None:
+    transform = Delaunay()
+
+    assert str(transform) == 'Delaunay()'
+    assert_one_point(transform)
+    assert_two_points(transform)
+    assert_three_points(transform)
+    assert_four_points(transform)
+
+
+@withPackage('torch_delaunay')
+def test_shull_delaunay() -> None:
+    transform = Delaunay()
+
+    assert str(transform) == 'Delaunay()'
+    assert_one_point(transform)
+    assert_two_points(transform)
+    assert_three_points(transform)
+    assert_four_points(transform)
diff --git a/test/utils/test_subgraph.py b/test/utils/test_subgraph.py
index 85161c55d..462bcb189 100644
--- a/test/utils/test_subgraph.py
+++ b/test/utils/test_subgraph.py
@@ -101,9 +101,12 @@ def test_k_hop_subgraph():
         [0, 1, 2, 3, 4, 5],
         [2, 2, 4, 4, 6, 6],
     ])
-
     subset, edge_index, mapping, edge_mask = k_hop_subgraph(
-        6, 2, edge_index, relabel_nodes=True)
+        node_idx=6,
+        num_hops=2,
+        edge_index=edge_index,
+        relabel_nodes=True,
+    )
     assert subset.tolist() == [2, 3, 4, 5, 6]
     assert edge_index.tolist() == [[0, 1, 2, 3], [2, 2, 4, 4]]
     assert mapping.tolist() == [4]
@@ -113,12 +116,29 @@ def test_k_hop_subgraph():
         [1, 2, 4, 5],
         [0, 1, 5, 6],
     ])
-
-    subset, edge_index, mapping, edge_mask = k_hop_subgraph([0, 6], 2,
-                                                            edge_index,
-                                                            relabel_nodes=True)
-
+    subset, edge_index, mapping, edge_mask = k_hop_subgraph(
+        node_idx=[0, 6],
+        num_hops=2,
+        edge_index=edge_index,
+        relabel_nodes=True,
+    )
     assert subset.tolist() == [0, 1, 2, 4, 5, 6]
     assert edge_index.tolist() == [[1, 2, 3, 4], [0, 1, 4, 5]]
     assert mapping.tolist() == [0, 5]
     assert edge_mask.tolist() == [True, True, True, True]
+
+    edge_index = torch.tensor([
+        [0, 1, 2, 3, 4, 4, 5],
+        [2, 2, 4, 4, 2, 6, 6],
+    ])
+    subset, edge_index, mapping, edge_mask = k_hop_subgraph(
+        node_idx=6,
+        num_hops=2,
+        edge_index=edge_index,
+        relabel_nodes=False,
+        directed=True,
+    )
+    assert subset.tolist() == [2, 3, 4, 5, 6]
+    assert edge_index.tolist() == [[2, 3, 4, 5], [4, 4, 6, 6]]
+    assert mapping.tolist() == [4]
+    assert edge_mask.tolist() == [False, False, True, True, False, True, True]