Allow floating-point slicing in Dataset (#7915)

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for floating-point slicing in `Dataset`, *e.g.*, `dataset[:0.9]` ([#7915](https://github.com/pyg-team/pytorch_geometric/pull/7915))
 - Added nightly GPU tests ([#7895](https://github.com/pyg-team/pytorch_geometric/pull/7895))
 - Added the `HalfHop` graph upsampling augmentation ([#7827](https://github.com/pyg-team/pytorch_geometric/pull/7827))
 - Added the `Wikidata5M` dataset ([#7864](https://github.com/pyg-team/pytorch_geometric/pull/7864))

test/datasets/test_enzymes.py

@@ -18,6 +18,7 @@ def test_enzymes(get_dataset):
     assert len(dataset.shuffle()) == 600
     assert len(dataset.shuffle(return_perm=True)) == 2
     assert len(dataset[:100]) == 100
+    assert len(dataset[0.1:0.2]) == 60
     assert len(dataset[torch.arange(100, dtype=torch.long)]) == 100
     mask = torch.zeros(600, dtype=torch.bool)
     mask[:100] = 1

torch_geometric/data/dataset.py

@@ -275,6 +275,14 @@ def index_select(self, idx: IndexType) -> 'Dataset':
         indices = self.indices()
 
         if isinstance(idx, slice):
+            start, stop, step = idx.start, idx.stop, idx.step
+            # Allow floating-point slicing, e.g., dataset[:0.9]
+            if isinstance(start, float):
+                start = round(start * len(self))
+            if isinstance(stop, float):
+                stop = round(stop * len(self))
+            idx = slice(start, stop, step)
+
             indices = indices[idx]
 
         elif isinstance(idx, Tensor) and idx.dtype == torch.long: