pyg-nightly 2.6.0.dev20240319__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/data/lightning/datamodule.py

@@ -13,7 +13,7 @@ from torch_geometric.typing import InputEdges, InputNodes, OptTensor
 try:
     from pytorch_lightning import LightningDataModule as PLLightningDataModule
     no_pytorch_lightning = False
-except (ImportError, ModuleNotFoundError):
+except ImportError:
     PLLightningDataModule = object  # type: ignore
     no_pytorch_lightning = True
 
@@ -221,7 +221,7 @@ class LightningDataset(LightningDataModule):
         speed.html>`__ are supported in order to correctly share data across
         all devices/processes:
 
-        .. code-block::
+        .. code-block:: python
 
             import pytorch_lightning as pl
             trainer = pl.Trainer(strategy="ddp_spawn", accelerator="gpu",
@@ -332,7 +332,7 @@ class LightningNodeData(LightningData):
         speed.html>`__ are supported in order to correctly share data across
         all devices/processes:
 
-        .. code-block::
+        .. code-block:: python
 
             import pytorch_lightning as pl
             trainer = pl.Trainer(strategy="ddp_spawn", accelerator="gpu",
@@ -525,7 +525,7 @@ class LightningLinkData(LightningData):
         speed.html>`__ are supported in order to correctly share data across
         all devices/processes:
 
-        .. code-block::
+        .. code-block:: python
 
             import pytorch_lightning as pl
             trainer = pl.Trainer(strategy="ddp_spawn", accelerator="gpu",

torch_geometric/data/separate.py

@@ -3,7 +3,7 @@ from typing import Any, Type, TypeVar
 
 from torch import Tensor
 
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.data.data import BaseData
 from torch_geometric.data.storage import BaseStorage
 from torch_geometric.typing import SparseTensor, TensorFrame
@@ -76,6 +76,11 @@ def _separate(
         value = narrow(values, cat_dim or 0, start, end - start)
         value = value.squeeze(0) if cat_dim is None else value
 
+        if isinstance(values, Index) and values._cat_metadata is not None:
+            # Reconstruct original `Index` metadata:
+            value._dim_size = values._cat_metadata.dim_size[idx]
+            value._is_sorted = values._cat_metadata.is_sorted[idx]
+
         if isinstance(values, EdgeIndex) and values._cat_metadata is not None:
             # Reconstruct original `EdgeIndex` metadata:
             value._sparse_size = values._cat_metadata.sparse_size[idx]

torch_geometric/data/storage.py

@@ -370,18 +370,20 @@ class BaseStorage(MutableMapping):
         self,
         start_time: Union[float, int],
         end_time: Union[float, int],
+        attr: str = 'time',
     ) -> Self:
-        if 'time' in self:
-            mask = (self.time >= start_time) & (self.time <= end_time)
+        if attr in self:
+            time = self[attr]
+            mask = (time >= start_time) & (time <= end_time)
 
-            if self.is_node_attr('time'):
+            if self.is_node_attr(attr):
                 keys = self.node_attrs()
-            elif self.is_edge_attr('time'):
+            elif self.is_edge_attr(attr):
                 keys = self.edge_attrs()
 
             self._select(keys, mask)
 
-            if self.is_node_attr('time') and 'num_nodes' in self:
+            if self.is_node_attr(attr) and 'num_nodes' in self:
                 self.num_nodes: Optional[int] = int(mask.sum())
 
         return self
@@ -443,9 +445,9 @@ class NodeStorage(BaseStorage):
                 return self.edge_index.sparse_size(0)
             if self.edge_index.sparse_size(1) is not None:
                 return self.edge_index.sparse_size(1)
-        if 'adj' in self and isinstance(self.adj, SparseTensor):
+        if 'adj' in self and isinstance(self.adj, (Tensor, SparseTensor)):
             return self.adj.size(0)
-        if 'adj_t' in self and isinstance(self.adj_t, SparseTensor):
+        if 'adj_t' in self and isinstance(self.adj_t, (Tensor, SparseTensor)):
             return self.adj_t.size(1)
         warnings.warn(
             f"Unable to accurately infer 'num_nodes' from the attribute set "
@@ -804,6 +806,10 @@ class GlobalStorage(NodeStorage, EdgeStorage):
             return False
 
         cat_dim = self._parent().__cat_dim__(key, value, self)
+
+        if not isinstance(cat_dim, int):
+            return False
+
         num_nodes, num_edges = self.num_nodes, self.num_edges
 
         if value.shape[cat_dim] != num_nodes:
@@ -850,6 +856,10 @@ class GlobalStorage(NodeStorage, EdgeStorage):
             return False
 
         cat_dim = self._parent().__cat_dim__(key, value, self)
+
+        if not isinstance(cat_dim, int):
+            return False
+
         num_nodes, num_edges = self.num_nodes, self.num_edges
 
         if value.shape[cat_dim] != num_edges:

torch_geometric/data/summary.py

@@ -117,7 +117,14 @@ class Summary:
             num_edges_per_type=num_edges_per_type,
         )
 
-    def __repr__(self) -> str:
+    def format(self, fmt: str = "psql") -> str:
+        r"""Formats summary statistics of the dataset.
+
+        Args:
+            fmt (str, optional): Summary tables format. Available table formats
+                can be found `here <https://github.com/astanin/python-tabulate?
+                tab=readme-ov-file#table-format>`__. (default: :obj:`"psql"`)
+        """
         from tabulate import tabulate
 
         body = f'{self.name} (#graphs={self.num_graphs}):\n'
@@ -127,7 +134,7 @@ class Summary:
         for field in Stats.__dataclass_fields__:
             row = [field] + [f'{getattr(s, field):.1f}' for s in stats]
             content.append(row)
-        body += tabulate(content, headers='firstrow', tablefmt='psql')
+        body += tabulate(content, headers='firstrow', tablefmt=fmt)
 
         if self.num_nodes_per_type is not None:
             content = [['']]
@@ -140,7 +147,7 @@ class Summary:
                 ]
                 content.append(row)
             body += "\nNumber of nodes per node type:\n"
-            body += tabulate(content, headers='firstrow', tablefmt='psql')
+            body += tabulate(content, headers='firstrow', tablefmt=fmt)
 
         if self.num_edges_per_type is not None:
             content = [['']]
@@ -156,6 +163,9 @@ class Summary:
                 ]
                 content.append(row)
             body += "\nNumber of edges per edge type:\n"
-            body += tabulate(content, headers='firstrow', tablefmt='psql')
+            body += tabulate(content, headers='firstrow', tablefmt=fmt)
 
         return body
+
+    def __repr__(self) -> str:
+        return self.format()

torch_geometric/data/temporal.py

@@ -156,8 +156,7 @@ class TemporalData(BaseData):
         return self.num_events
 
     def __call__(self, *args: List[str]) -> Iterable:
-        for key, value in self._store.items(*args):
-            yield key, value
+        yield from self._store.items(*args)
 
     def __copy__(self):
         out = self.__class__.__new__(self.__class__)

torch_geometric/datasets/__init__.py

@@ -61,7 +61,6 @@ from .gemsec import GemsecDeezer
 from .twitch import Twitch
 from .airports import Airports
 from .lrgb import LRGBDataset
-from .neurograph import NeuroGraphDataset
 from .malnet_tiny import MalNetTiny
 from .omdb import OMDB
 from .polblogs import PolBlogs
@@ -76,6 +75,11 @@ from .jodie import JODIEDataset
 from .wikidata import Wikidata5M
 from .myket import MyketDataset
 from .brca_tgca import BrcaTcga
+from .neurograph import NeuroGraphDataset
+from .web_qsp_dataset import WebQSPDataset
+from .git_mol_dataset import GitMolDataset
+from .molecule_gpt_dataset import MoleculeGPTDataset
+from .tag_dataset import TAGDataset
 
 from .dbp15k import DBP15K
 from .aminer import AMiner
@@ -93,6 +97,9 @@ from .amazon_book import AmazonBook
 from .hm import HM
 from .ose_gvcs import OSE_GVCS
 from .rcdd import RCDD
+from .opf import OPFDataset
+
+from .cornell import CornellTemporalHyperGraphDataset
 
 from .fake import FakeDataset, FakeHeteroDataset
 from .sbm_dataset import StochasticBlockModelDataset
@@ -185,6 +192,10 @@ homo_datasets = [
     'MyketDataset',
     'BrcaTcga',
     'NeuroGraphDataset',
+    'WebQSPDataset',
+    'GitMolDataset',
+    'MoleculeGPTDataset',
+    'TAGDataset',
 ]
 
 hetero_datasets = [
@@ -204,6 +215,10 @@ hetero_datasets = [
     'HM',
     'OSE_GVCS',
     'RCDD',
+    'OPFDataset',
+]
+hyper_datasets = [
+    'CornellTemporalHyperGraphDataset',
 ]
 synthetic_datasets = [
     'FakeDataset',
@@ -218,4 +233,4 @@ synthetic_datasets = [
     'BAShapes',
 ]
 
-__all__ = homo_datasets + hetero_datasets + synthetic_datasets
+__all__ = homo_datasets + hetero_datasets + hyper_datasets + synthetic_datasets

torch_geometric/datasets/actor.py

@@ -19,17 +19,15 @@ class Actor(InMemoryDataset):
     actor's Wikipedia.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in an
             :obj:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
-            transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before being saved to
+            disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 
@@ -76,7 +74,7 @@ class Actor(InMemoryDataset):
             download_url(f'{self.url}/splits/{f}', self.raw_dir)
 
     def process(self) -> None:
-        with open(self.raw_paths[0], 'r') as f:
+        with open(self.raw_paths[0]) as f:
             node_data = [x.split('\t') for x in f.read().split('\n')[1:-1]]
 
             rows, cols = [], []
@@ -93,7 +91,7 @@ class Actor(InMemoryDataset):
             for n_id, _, label in node_data:
                 y[int(n_id)] = int(label)
 
-        with open(self.raw_paths[1], 'r') as f:
+        with open(self.raw_paths[1]) as f:
             edge_data = f.read().split('\n')[1:-1]
             edge_indices = [[int(v) for v in r.split('\t')] for r in edge_data]
             edge_index = torch.tensor(edge_indices).t().contiguous()

torch_geometric/datasets/airfrans.py

@@ -2,14 +2,13 @@ import json
 import os
 from typing import Callable, List, Optional
 
-import torch
-
 from torch_geometric.data import (
     Data,
     InMemoryDataset,
     download_url,
     extract_zip,
 )
+from torch_geometric.io import fs
 
 
 class AirfRANS(InMemoryDataset):
@@ -47,26 +46,24 @@ class AirfRANS(InMemoryDataset):
         :obj:`torch_geometric.transforms.RadiusGraph` transform.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        task (str): The task to study (:obj:`"full"`, :obj:`"scarce"`,
+        root: Root directory where the dataset should be saved.
+        task: The task to study (:obj:`"full"`, :obj:`"scarce"`,
             :obj:`"reynolds"`, :obj:`"aoa"`) that defines the utilized training
             and test splits.
-        train (bool, optional): If :obj:`True`, loads the training dataset,
-            otherwise the test dataset. (default: :obj:`True`)
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        train: If :obj:`True`, loads the training dataset, otherwise the test
+            dataset.
+        transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        pre_filter (callable, optional): A function that takes in an
+            being saved to disk.
+        pre_filter: A function that takes in an
             :obj:`torch_geometric.data.Data` object and returns a boolean
             value, indicating whether the data object should be included in the
-            final dataset. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            final dataset.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 
@@ -123,13 +120,13 @@ class AirfRANS(InMemoryDataset):
         os.unlink(path)
 
     def process(self) -> None:
-        with open(self.raw_paths[1], 'r') as f:
+        with open(self.raw_paths[1]) as f:
             manifest = json.load(f)
         total = manifest['full_train'] + manifest['full_test']
         partial = set(manifest[f'{self.task}_{self.split}'])
 
         data_list = []
-        raw_data = torch.load(self.raw_paths[0])
+        raw_data = fs.torch_load(self.raw_paths[0])
         for k, s in enumerate(total):
             if s in partial:
                 data = Data(**raw_data[k])

torch_geometric/datasets/airports.py

@@ -14,22 +14,20 @@ class Airports(InMemoryDataset):
     and labels correspond to activity levels.
     Features are given by one-hot encoded node identifiers, as described in the
     `"GraLSP: Graph Neural Networks with Local Structural Patterns"
-    ` <https://arxiv.org/abs/1911.07675>`_ paper.
+    <https://arxiv.org/abs/1911.07675>`_ paper.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        name (str): The name of the dataset (:obj:`"USA"`, :obj:`"Brazil"`,
+        root: Root directory where the dataset should be saved.
+        name: The name of the dataset (:obj:`"USA"`, :obj:`"Brazil"`,
             :obj:`"Europe"`).
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
         pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
     """
     edge_url = ('https://github.com/leoribeiro/struc2vec/'
                 'raw/master/graph/{}-airports.edgelist')
@@ -75,7 +73,7 @@ class Airports(InMemoryDataset):
 
     def process(self) -> None:
         index_map, ys = {}, []
-        with open(self.raw_paths[1], 'r') as f:
+        with open(self.raw_paths[1]) as f:
             rows = f.read().split('\n')[1:-1]
             for i, row in enumerate(rows):
                 idx, label = row.split()
@@ -85,7 +83,7 @@ class Airports(InMemoryDataset):
         x = torch.eye(y.size(0))
 
         edge_indices = []
-        with open(self.raw_paths[0], 'r') as f:
+        with open(self.raw_paths[0]) as f:
             rows = f.read().split('\n')[:-1]
             for row in rows:
                 src, dst = row.split()

torch_geometric/datasets/amazon.py

@@ -15,19 +15,16 @@ class Amazon(InMemoryDataset):
     map goods to their respective product category.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        name (str): The name of the dataset (:obj:`"Computers"`,
-            :obj:`"Photo"`).
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        root: Root directory where the dataset should be saved.
+        name: The name of the dataset (:obj:`"Computers"`, :obj:`"Photo"`).
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

torch_geometric/datasets/amazon_book.py

@@ -14,17 +14,16 @@ class AmazonBook(InMemoryDataset):
     No labels or features are provided.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.HeteroData` object and returns a
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in an
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            every access. (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.HeteroData` object and returns a
+            every access.
+        pre_transform: A function/transform that takes in an
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
     """
     url = ('https://raw.githubusercontent.com/gusye1234/LightGCN-PyTorch/'
            'master/data/amazon-book')
@@ -67,7 +66,7 @@ class AmazonBook(InMemoryDataset):
         attr_names = ['edge_index', 'edge_label_index']
         for path, attr_name in zip(self.raw_paths[2:], attr_names):
             rows, cols = [], []
-            with open(path, 'r') as f:
+            with open(path) as f:
                 lines = f.readlines()
             for line in lines:
                 indices = line.strip().split(' ')

torch_geometric/datasets/amazon_products.py

@@ -3,7 +3,6 @@ import os.path as osp
 from typing import Callable, List, Optional
 
 import numpy as np
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import Data, InMemoryDataset, download_google_url
@@ -15,17 +14,15 @@ class AmazonProducts(InMemoryDataset):
     containing products and its categories.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in an
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 
@@ -73,6 +70,8 @@ class AmazonProducts(InMemoryDataset):
         download_google_url(self.role_id, self.raw_dir, 'role.json')
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         f = np.load(osp.join(self.raw_dir, 'adj_full.npz'))
         adj = sp.csr_matrix((f['data'], f['indices'], f['indptr']), f['shape'])
         adj = adj.tocoo()

torch_geometric/datasets/aminer.py

@@ -24,17 +24,16 @@ class AMiner(InMemoryDataset):
     truth labels for a subset of nodes.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.HeteroData` object and returns a
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            every access. (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.HeteroData` object and returns a
+            every access.
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.HeteroData` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
     """
 
     url = 'https://www.dropbox.com/s/1bnz8r7mofx0osf/net_aminer.zip?dl=1'

torch_geometric/datasets/aqsol.py

@@ -30,25 +30,22 @@ class AQSOL(InMemoryDataset):
     the :class:`~torch_geometric.datasets.ZINC` dataset.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        split (str, optional): If :obj:`"train"`, loads the training dataset.
+        root: Root directory where the dataset should be saved.
+        split: If :obj:`"train"`, loads the training dataset.
             If :obj:`"val"`, loads the validation dataset.
             If :obj:`"test"`, loads the test dataset.
-            (default: :obj:`"train"`)
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
+            being saved to disk.
         pre_filter (callable, optional): A function that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a boolean
+            :class:`torch_geometric.data.Data` object and returns a boolean
             value, indicating whether the data object should be included in
-            the final dataset. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            the final dataset.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 

torch_geometric/datasets/attributed_graph_dataset.py

@@ -2,7 +2,6 @@ import os
 import os.path as osp
 from typing import Callable, List, Optional
 
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import (
@@ -20,21 +19,19 @@ class AttributedGraphDataset(InMemoryDataset):
     <https://arxiv.org/abs/2009.00826>`_ paper.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        name (str): The name of the dataset (:obj:`"Wiki"`, :obj:`"Cora"`
+        root: Root directory where the dataset should be saved.
+        name: The name of the dataset (:obj:`"Wiki"`, :obj:`"Cora"`,
             :obj:`"CiteSeer"`, :obj:`"PubMed"`, :obj:`"BlogCatalog"`,
             :obj:`"PPI"`, :obj:`"Flickr"`, :obj:`"Facebook"`, :obj:`"Twitter"`,
             :obj:`"TWeibo"`, :obj:`"MAG"`).
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            being saved to disk.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**
 
@@ -156,6 +153,7 @@ class AttributedGraphDataset(InMemoryDataset):
 
     def process(self) -> None:
         import pandas as pd
+        import scipy.sparse as sp
 
         x = sp.load_npz(self.raw_paths[0]).tocsr()
         if x.shape[-1] > 10000 or self.name == 'mag':
@@ -172,7 +170,7 @@ class AttributedGraphDataset(InMemoryDataset):
                          engine='python')
         edge_index = torch.from_numpy(df.values).t().contiguous()
 
-        with open(self.raw_paths[2], 'r') as f:
+        with open(self.raw_paths[2]) as f:
             rows = f.read().split('\n')[:-1]
             ys = [[int(y) - 1 for y in row.split()[1:]] for row in rows]
             multilabel = max([len(y) for y in ys]) > 1

torch_geometric/datasets/ba_multi_shapes.py

@@ -25,21 +25,19 @@ class BAMultiShapesDataset(InMemoryDataset):
     This dataset is pre-computed from the official implementation.
 
     Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a transformed
+        root: Root directory where the dataset should be saved.
+        transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a transformed
             version. The data object will be transformed before every access.
-            (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.Data` object and returns a
+        pre_transform: A function/transform that takes in a
+            :class:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        pre_filter (callable, optional): A function that takes in an
-            :obj:`torch_geometric.data.Data` object and returns a boolean
+            being saved to disk.
+        pre_filter: A function that takes in a
+            :class:`torch_geometric.data.Data` object and returns a boolean
             value, indicating whether the data object should be included in the
-            final dataset. (default: :obj:`None`)
-        force_reload (bool, optional): Whether to re-process the dataset.
-            (default: :obj:`False`)
+            final dataset.
+        force_reload: Whether to re-process the dataset.
 
     **STATS:**