pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/data/database.py

@@ -1,4 +1,4 @@
-import pickle
+import io
 import warnings
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
@@ -111,13 +111,17 @@ class Database(ABC):
             for key, value in schema_dict.items()
         }
 
+    @abstractmethod
     def connect(self) -> None:
         r"""Connects to the database.
         Databases will automatically connect on instantiation.
         """
+        raise NotImplementedError
 
+    @abstractmethod
     def close(self) -> None:
         r"""Closes the connection to the database."""
+        raise NotImplementedError
 
     @abstractmethod
     def insert(self, index: int, data: Any) -> None:
@@ -496,7 +500,9 @@ class SQLiteDatabase(Database):
                 out.append(col)
 
             else:
-                out.append(pickle.dumps(col))
+                buffer = io.BytesIO()
+                torch.save(col, buffer)
+                out.append(buffer.getvalue())
 
         return out
 
@@ -559,7 +565,10 @@ class SQLiteDatabase(Database):
                 out_dict[key] = value
 
             else:
-                out_dict[key] = pickle.loads(value)
+                out_dict[key] = torch.load(
+                    io.BytesIO(value),
+                    weights_only=False,
+                )
 
         # In case `0` exists as integer in the schema, this means that the
         # schema was passed as either a single entry or a tuple:
@@ -644,7 +653,12 @@ class RocksDatabase(Database):
         # Ensure that data is not a view of a larger tensor:
         if isinstance(row, Tensor):
             row = row.clone()
-        return pickle.dumps(row)
+        buffer = io.BytesIO()
+        torch.save(row, buffer)
+        return buffer.getvalue()
 
     def _deserialize(self, row: bytes) -> Any:
-        return pickle.loads(row)
+        return torch.load(
+            io.BytesIO(row),
+            weights_only=False,
+        )

torch_geometric/data/dataset.py

@@ -166,10 +166,11 @@ class Dataset(torch.utils.data.Dataset):
         elif y.numel() == y.size(0) and torch.is_floating_point(y):
             num_classes = torch.unique(y).numel()
             if num_classes > 2:
-                warnings.warn("Found floating-point labels while calling "
-                              "`dataset.num_classes`. Returning the number of "
-                              "unique elements. Please make sure that this "
-                              "is expected before proceeding.")
+                warnings.warn(
+                    "Found floating-point labels while calling "
+                    "`dataset.num_classes`. Returning the number of "
+                    "unique elements. Please make sure that this "
+                    "is expected before proceeding.", stacklevel=2)
             return num_classes
         else:
             return y.size(-1)
@@ -235,20 +236,24 @@ class Dataset(torch.utils.data.Dataset):
 
     def _process(self):
         f = osp.join(self.processed_dir, 'pre_transform.pt')
-        if osp.exists(f) and torch.load(f) != _repr(self.pre_transform):
+        if not self.force_reload and osp.exists(f) and torch.load(
+                f, weights_only=False) != _repr(self.pre_transform):
             warnings.warn(
                 "The `pre_transform` argument differs from the one used in "
                 "the pre-processed version of this dataset. If you want to "
                 "make use of another pre-processing technique, pass "
-                "`force_reload=True` explicitly to reload the dataset.")
+                "`force_reload=True` explicitly to reload the dataset.",
+                stacklevel=2)
 
         f = osp.join(self.processed_dir, 'pre_filter.pt')
-        if osp.exists(f) and torch.load(f) != _repr(self.pre_filter):
+        if not self.force_reload and osp.exists(f) and torch.load(
+                f, weights_only=False) != _repr(self.pre_filter):
             warnings.warn(
                 "The `pre_filter` argument differs from the one used in "
                 "the pre-processed version of this dataset. If you want to "
                 "make use of another pre-fitering technique, pass "
-                "`force_reload=True` explicitly to reload the dataset.")
+                "`force_reload=True` explicitly to reload the dataset.",
+                stacklevel=2)
 
         if not self.force_reload and files_exist(self.processed_paths):
             return
@@ -381,7 +386,7 @@ class Dataset(torch.utils.data.Dataset):
         r"""Converts the dataset into a :class:`torch.utils.data.DataPipe`.
 
         The returned instance can then be used with :pyg:`PyG's` built-in
-        :class:`DataPipes` for baching graphs as follows:
+        :class:`DataPipes` for batching graphs as follows:
 
         .. code-block:: python
 

torch_geometric/data/extract.py

@@ -28,7 +28,7 @@ def extract_tar(
     """
     maybe_log(path, log)
     with tarfile.open(path, mode) as f:
-        f.extractall(folder)
+        f.extractall(folder, filter='data')
 
 
 def extract_zip(path: str, folder: str, log: bool = True) -> None:

torch_geometric/data/feature_store.py

@@ -11,7 +11,7 @@ This particular feature store abstraction makes a few key assumptions:
 * A feature can be uniquely identified from any associated attributes specified
   in `TensorAttr`.
 
-It is the job of a feature store implementor class to handle these assumptions
+It is the job of a feature store implementer class to handle these assumptions
 properly. For example, a simple in-memory feature store implementation may
 concatenate all metadata values with a feature index and use this as a unique
 index in a KV store. More complicated implementations may choose to partition
@@ -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__)
@@ -358,7 +352,7 @@ class FeatureStore(ABC):
 
         .. note::
             The default implementation simply iterates over all calls to
-            :meth:`get_tensor`. Implementor classes that can provide
+            :meth:`get_tensor`. Implementer classes that can provide
             additional, more performant functionality are recommended to
             to override this method.
 
@@ -415,10 +409,10 @@ class FeatureStore(ABC):
     def update_tensor(self, tensor: FeatureTensorType, *args,
                       **kwargs) -> bool:
         r"""Updates a :obj:`tensor` in the :class:`FeatureStore` with a new
-        value. Returns whether the update was succesful.
+        value. Returns whether the update was successful.
 
         .. note::
-            Implementor classes can choose to define more efficient update
+            Implementer classes can choose to define more efficient update
             methods; the default performs a removal and insertion.
 
         Args:
@@ -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

torch_geometric/data/graph_store.py

@@ -10,7 +10,7 @@ This particular graph store abstraction makes a few key assumptions:
   support dynamic modification of edge indices once they have been inserted
   into the graph store.
 
-It is the job of a graph store implementor class to handle these assumptions
+It is the job of a graph store implementer class to handle these assumptions
 properly. For example, a simple in-memory graph store implementation may
 concatenate all metadata values with an edge index and use this as a unique
 index in a KV store. More complicated implementations may choose to partition
@@ -261,7 +261,8 @@ class GraphStore(ABC):
                 col = ptr2index(col)
 
             if attr.layout != EdgeLayout.CSR:  # COO->CSR
-                num_rows = attr.size[0] if attr.size else int(row.max()) + 1
+                num_rows = attr.size[0] if attr.size is not None else int(
+                    row.max()) + 1
                 row, perm = index_sort(row, max_value=num_rows)
                 col = col[perm]
                 row = index2ptr(row, num_rows)

torch_geometric/data/hetero_data.py

@@ -282,6 +282,21 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         r"""Returns a list of edge type and edge storage pairs."""
         return list(self._edge_store_dict.items())
 
+    @property
+    def input_type(self) -> Optional[Union[NodeType, EdgeType]]:
+        r"""Returns the seed/input node/edge type of the graph in case it
+        refers to a sampled subgraph, *e.g.*, obtained via
+        :class:`~torch_geometric.loader.NeighborLoader` or
+        :class:`~torch_geometric.loader.LinkNeighborLoader`.
+        """
+        for node_type, store in self.node_items():
+            if hasattr(store, 'input_id'):
+                return node_type
+        for edge_type, store in self.edge_items():
+            if hasattr(store, 'input_id'):
+                return edge_type
+        return None
+
     def to_dict(self) -> Dict[str, Any]:
         out_dict: Dict[str, Any] = {}
         out_dict['_global_store'] = self._global_store.to_dict()
@@ -472,6 +487,77 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
 
         return status
 
+    def connected_components(self) -> List[Self]:
+        r"""Extracts connected components of the heterogeneous graph using
+        a union-find algorithm. The components are returned as a list of
+        :class:`~torch_geometric.data.HeteroData` objects.
+
+        .. code-block::
+
+            data = HeteroData()
+            data["red"].x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
+            data["blue"].x = torch.tensor([[5.0], [6.0]])
+            data["red", "to", "red"].edge_index = torch.tensor(
+                [[0, 1, 2, 3], [1, 0, 3, 2]], dtype=torch.long
+            )
+
+            components = data.connected_components()
+            print(len(components))
+            >>> 4
+
+            print(components[0])
+            >>> HeteroData(
+                red={x: tensor([[1.], [2.]])},
+                blue={x: tensor([[]])},
+                red, to, red={edge_index: tensor([[0, 1], [1, 0]])}
+            )
+
+        Returns:
+            List[HeteroData]: A list of connected components.
+        """
+        # Initialize union-find structures
+        self._parents: Dict[Tuple[str, int], Tuple[str, int]] = {}
+        self._ranks: Dict[Tuple[str, int], int] = {}
+
+        # Union-Find algorithm to find connected components
+        for edge_type in self.edge_types:
+            src, _, dst = edge_type
+            edge_index = self[edge_type].edge_index
+            for src_node, dst_node in edge_index.t().tolist():
+                self._union((src, src_node), (dst, dst_node))
+
+        # Rerun _find_parent to ensure all nodes are covered correctly
+        for node_type in self.node_types:
+            for node_index in range(self[node_type].num_nodes):
+                self._find_parent((node_type, node_index))
+
+        # Group nodes by their representative parent
+        components_map = defaultdict(list)
+        for node, parent in self._parents.items():
+            components_map[parent].append(node)
+        del self._parents
+        del self._ranks
+
+        components: List[Self] = []
+        for nodes in components_map.values():
+            # Prefill subset_dict with all node types to ensure all are present
+            subset_dict = {node_type: [] for node_type in self.node_types}
+
+            # Convert the list of (node_type, node_id) tuples to a subset_dict
+            for node_type, node_id in nodes:
+                subset_dict[node_type].append(node_id)
+
+            # Convert lists to tensors
+            for node_type, node_ids in subset_dict.items():
+                subset_dict[node_type] = torch.tensor(node_ids,
+                                                      dtype=torch.long)
+
+            # Use the existing subgraph function to do all the heavy lifting
+            component_data = self.subgraph(subset_dict)
+            components.append(component_data)
+
+        return components
+
     def debug(self):
         pass  # TODO
 
@@ -551,7 +637,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
             This is equivalent to writing :obj:`data.x_dict`.
 
         Args:
-            key (str): The attribute to collect from all node and ege types.
+            key (str): The attribute to collect from all node and edge types.
             allow_empty (bool, optional): If set to :obj:`True`, will not raise
                 an error in case the attribute does not exit in any node or
                 edge type. (default: :obj:`False`)
@@ -570,12 +656,13 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         global _DISPLAYED_TYPE_NAME_WARNING
         if not _DISPLAYED_TYPE_NAME_WARNING and '__' in name:
             _DISPLAYED_TYPE_NAME_WARNING = True
-            warnings.warn(f"There exist type names in the "
-                          f"'{self.__class__.__name__}' object that contain "
-                          f"double underscores '__' (e.g., '{name}'). This "
-                          f"may lead to unexpected behavior. To avoid any "
-                          f"issues, ensure that your type names only contain "
-                          f"single underscores.")
+            warnings.warn(
+                f"There exist type names in the "
+                f"'{self.__class__.__name__}' object that contain "
+                f"double underscores '__' (e.g., '{name}'). This "
+                f"may lead to unexpected behavior. To avoid any "
+                f"issues, ensure that your type names only contain "
+                f"single underscores.", stacklevel=2)
 
     def get_node_store(self, key: NodeType) -> NodeStorage:
         r"""Gets the :class:`~torch_geometric.data.storage.NodeStorage` object
@@ -1132,6 +1219,51 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
 
         return list(edge_attrs.values())
 
+    # Connected Components Helper Functions ###################################
+
+    def _find_parent(self, node: Tuple[str, int]) -> Tuple[str, int]:
+        r"""Finds and returns the representative parent of the given node in a
+        disjoint-set (union-find) data structure. Implements path compression
+        to optimize future queries.
+
+        Args:
+            node (tuple[str, int]): The node for which to find the parent.
+            First element is the node type, second is the node index.
+
+        Returns:
+            tuple[str, int]: The representative parent of the node.
+        """
+        if node not in self._parents:
+            self._parents[node] = node
+            self._ranks[node] = 0
+        if self._parents[node] != node:
+            self._parents[node] = self._find_parent(self._parents[node])
+        return self._parents[node]
+
+    def _union(self, node1: Tuple[str, int], node2: Tuple[str, int]):
+        r"""Merges the node1 and node2 in the disjoint-set data structure.
+
+        Finds the root parents of node1 and node2 using the _find_parent
+        method. If they belong to different sets, updates the parent of
+        root2 to be root1, effectively merging the two sets.
+
+        Args:
+            node1 (Tuple[str, int]): The first node to union. First element is
+                the node type, second is the node index.
+            node2 (Tuple[str, int]): The second node to union. First element is
+                the node type, second is the node index.
+        """
+        root1 = self._find_parent(node1)
+        root2 = self._find_parent(node2)
+        if root1 != root2:
+            if self._ranks[root1] < self._ranks[root2]:
+                self._parents[root1] = root2
+            elif self._ranks[root1] > self._ranks[root2]:
+                self._parents[root2] = root1
+            else:
+                self._parents[root2] = root1
+                self._ranks[root1] += 1
+
 
 # Helper functions ############################################################
 

torch_geometric/data/hypergraph_data.py

@@ -39,7 +39,7 @@ class HyperGraphData(Data):
         edge_index (LongTensor, optional): Hyperedge tensor
             with shape :obj:`[2, num_edges*num_nodes_per_edge]`.
             Where `edge_index[1]` denotes the hyperedge index and
-            `edge_index[0]` denotes the node indicies that are connected
+            `edge_index[0]` denotes the node indices that are connected
             by the hyperedge. (default: :obj:`None`)
             (default: :obj:`None`)
         edge_attr (torch.Tensor, optional): Edge feature matrix with shape
@@ -223,4 +223,4 @@ def warn_or_raise(msg: str, raise_on_error: bool = True) -> None:
     if raise_on_error:
         raise ValueError(msg)
     else:
-        warnings.warn(msg)
+        warnings.warn(msg, stacklevel=2)

torch_geometric/data/in_memory_dataset.py

@@ -297,7 +297,7 @@ class InMemoryDataset(Dataset):
             self._data_list = None
         msg += f' {msg4}'
 
-        warnings.warn(msg)
+        warnings.warn(msg, stacklevel=2)
 
         return self._data
 
@@ -346,7 +346,7 @@ class InMemoryDataset(Dataset):
 
 def nested_iter(node: Union[Mapping, Sequence]) -> Iterable:
     if isinstance(node, Mapping):
-        for key, value in node.items():
+        for value in node.values():
             yield from nested_iter(value)
     elif isinstance(node, Sequence):
         yield from enumerate(node)

torch_geometric/data/lightning/datamodule.py

@@ -11,21 +11,27 @@ from torch_geometric.sampler import BaseSampler, NeighborSampler
 from torch_geometric.typing import InputEdges, InputNodes, OptTensor
 
 try:
-    from pytorch_lightning import LightningDataModule as PLLightningDataModule
-    no_pytorch_lightning = False
+    from lightning.pytorch import LightningDataModule as _LightningDataModule
+    _pl_is_available = True
 except ImportError:
-    PLLightningDataModule = object  # type: ignore
-    no_pytorch_lightning = True
+    try:
+        from pytorch_lightning import \
+            LightningDataModule as _LightningDataModule
+        _pl_is_available = True
+    except ImportError:
+        _pl_is_available = False
+        _LightningDataModule = object
 
 
-class LightningDataModule(PLLightningDataModule):
+class LightningDataModule(_LightningDataModule):
     def __init__(self, has_val: bool, has_test: bool, **kwargs: Any) -> None:
         super().__init__()
 
-        if no_pytorch_lightning:
+        if not _pl_is_available:
             raise ModuleNotFoundError(
-                "No module named 'pytorch_lightning' found on this machine. "
-                "Run 'pip install pytorch_lightning' to install the library.")
+                "No module named 'pytorch_lightning' (or 'lightning') found "
+                "in your Python environment. Run 'pip install "
+                "pytorch_lightning' or 'pip install lightning'")
 
         if not has_val:
             self.val_dataloader = None  # type: ignore
@@ -40,9 +46,11 @@ class LightningDataModule(PLLightningDataModule):
                           kwargs.get('num_workers', 0) > 0)
 
         if 'shuffle' in kwargs:
-            warnings.warn(f"The 'shuffle={kwargs['shuffle']}' option is "
-                          f"ignored in '{self.__class__.__name__}'. Remove it "
-                          f"from the argument list to disable this warning")
+            warnings.warn(
+                f"The 'shuffle={kwargs['shuffle']}' option is "
+                f"ignored in '{self.__class__.__name__}'. Remove it "
+                f"from the argument list to disable this warning",
+                stacklevel=2)
             del kwargs['shuffle']
 
         self.kwargs = kwargs
@@ -74,34 +82,39 @@ class LightningData(LightningDataModule):
             raise ValueError(f"Undefined 'loader' option (got '{loader}')")
 
         if loader == 'full' and kwargs['batch_size'] != 1:
-            warnings.warn(f"Re-setting 'batch_size' to 1 in "
-                          f"'{self.__class__.__name__}' for loader='full' "
-                          f"(got '{kwargs['batch_size']}')")
+            warnings.warn(
+                f"Re-setting 'batch_size' to 1 in "
+                f"'{self.__class__.__name__}' for loader='full' "
+                f"(got '{kwargs['batch_size']}')", stacklevel=2)
             kwargs['batch_size'] = 1
 
         if loader == 'full' and kwargs['num_workers'] != 0:
-            warnings.warn(f"Re-setting 'num_workers' to 0 in "
-                          f"'{self.__class__.__name__}' for loader='full' "
-                          f"(got '{kwargs['num_workers']}')")
+            warnings.warn(
+                f"Re-setting 'num_workers' to 0 in "
+                f"'{self.__class__.__name__}' for loader='full' "
+                f"(got '{kwargs['num_workers']}')", stacklevel=2)
             kwargs['num_workers'] = 0
 
         if loader == 'full' and kwargs.get('sampler') is not None:
-            warnings.warn("'sampler' option is not supported for "
-                          "loader='full'")
+            warnings.warn(
+                "'sampler' option is not supported for "
+                "loader='full'", stacklevel=2)
             kwargs.pop('sampler', None)
 
         if loader == 'full' and kwargs.get('batch_sampler') is not None:
-            warnings.warn("'batch_sampler' option is not supported for "
-                          "loader='full'")
+            warnings.warn(
+                "'batch_sampler' option is not supported for "
+                "loader='full'", stacklevel=2)
             kwargs.pop('batch_sampler', None)
 
         super().__init__(has_val, has_test, **kwargs)
 
         if loader == 'full':
             if kwargs.get('pin_memory', False):
-                warnings.warn(f"Re-setting 'pin_memory' to 'False' in "
-                              f"'{self.__class__.__name__}' for loader='full' "
-                              f"(got 'True')")
+                warnings.warn(
+                    f"Re-setting 'pin_memory' to 'False' in "
+                    f"'{self.__class__.__name__}' for loader='full' "
+                    f"(got 'True')", stacklevel=2)
             self.kwargs['pin_memory'] = False
 
         self.data = data
@@ -127,10 +140,11 @@ class LightningData(LightningDataModule):
                 graph_sampler.__class__,
             )
             if len(sampler_kwargs) > 0:
-                warnings.warn(f"Ignoring the arguments "
-                              f"{list(sampler_kwargs.keys())} in "
-                              f"'{self.__class__.__name__}' since a custom "
-                              f"'graph_sampler' was passed")
+                warnings.warn(
+                    f"Ignoring the arguments "
+                    f"{list(sampler_kwargs.keys())} in "
+                    f"'{self.__class__.__name__}' since a custom "
+                    f"'graph_sampler' was passed", stacklevel=2)
             self.graph_sampler = graph_sampler
 
         else:

torch_geometric/data/storage.py

@@ -454,7 +454,7 @@ class NodeStorage(BaseStorage):
             f"'{set(self.keys())}'. Please explicitly set 'num_nodes' as an "
             f"attribute of " +
             ("'data'" if self._key is None else f"'data[{self._key}]'") +
-            " to suppress this warning")
+            " to suppress this warning", stacklevel=2)
         if 'edge_index' in self and isinstance(self.edge_index, Tensor):
             if self.edge_index.numel() > 0:
                 return int(self.edge_index.max()) + 1
@@ -806,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:
@@ -852,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/datasets/__init__.py

@@ -30,6 +30,7 @@ from .faust import FAUST
 from .dynamic_faust import DynamicFAUST
 from .shapenet import ShapeNet
 from .modelnet import ModelNet
+from .medshapenet import MedShapeNet
 from .coma import CoMA
 from .shrec2016 import SHREC2016
 from .tosca import TOSCA
@@ -61,7 +62,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 +76,15 @@ 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, CWQDataset
+from .git_mol_dataset import GitMolDataset
+from .molecule_gpt_dataset import MoleculeGPTDataset
+from .instruct_mol_dataset import InstructMolDataset
+from .protein_mpnn_dataset import ProteinMPNNDataset
+from .tag_dataset import TAGDataset
+from .city import CityNetwork
+from .teeth3ds import Teeth3DS
 
 from .dbp15k import DBP15K
 from .aminer import AMiner
@@ -141,6 +150,7 @@ homo_datasets = [
     'DynamicFAUST',
     'ShapeNet',
     'ModelNet',
+    'MedShapeNet',
     'CoMA',
     'SHREC2016',
     'TOSCA',
@@ -188,6 +198,15 @@ homo_datasets = [
     'MyketDataset',
     'BrcaTcga',
     'NeuroGraphDataset',
+    'WebQSPDataset',
+    'CWQDataset',
+    'GitMolDataset',
+    'MoleculeGPTDataset',
+    'InstructMolDataset',
+    'ProteinMPNNDataset',
+    'TAGDataset',
+    'CityNetwork',
+    'Teeth3DS',
 ]
 
 hetero_datasets = [