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

torch_geometric/data/dataset.py

@@ -235,7 +235,8 @@ 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 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 "
@@ -243,7 +244,8 @@ class Dataset(torch.utils.data.Dataset):
                 "`force_reload=True` explicitly to reload the dataset.")
 
         f = osp.join(self.processed_dir, 'pre_filter.pt')
-        if osp.exists(f) and torch.load(f) != _repr(self.pre_filter):
+        if 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 "
@@ -367,15 +369,21 @@ class Dataset(torch.utils.data.Dataset):
         from torch_geometric.data.summary import Summary
         return Summary.from_dataset(self)
 
-    def print_summary(self) -> None:
-        r"""Prints summary statistics of the dataset to the console."""
-        print(str(self.get_summary()))
+    def print_summary(self, fmt: str = "psql") -> None:
+        r"""Prints summary statistics of the dataset to the console.
+
+        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"`)
+        """
+        print(self.get_summary().format(fmt=fmt))
 
     def to_datapipe(self) -> Any:
         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/feature_store.py

@@ -28,6 +28,7 @@ from typing import Any, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
+from torch import Tensor
 
 from torch_geometric.typing import FeatureTensorType, NodeType
 from torch_geometric.utils.mixin import CastMixin
@@ -73,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`.
@@ -229,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__)
@@ -282,7 +277,6 @@ class FeatureStore(ABC):
     @abstractmethod
     def _put_tensor(self, tensor: FeatureTensorType, attr: TensorAttr) -> bool:
         r"""To be implemented by :class:`FeatureStore` subclasses."""
-        pass
 
     def put_tensor(self, tensor: FeatureTensorType, *args, **kwargs) -> bool:
         r"""Synchronously adds a :obj:`tensor` to the :class:`FeatureStore`.
@@ -308,7 +302,6 @@ class FeatureStore(ABC):
     @abstractmethod
     def _get_tensor(self, attr: TensorAttr) -> Optional[FeatureTensorType]:
         r"""To be implemented by :class:`FeatureStore` subclasses."""
-        pass
 
     def get_tensor(
         self,
@@ -329,8 +322,6 @@ class FeatureStore(ABC):
         Raises:
             ValueError: If the input :class:`TensorAttr` is not fully
                 specified.
-            KeyError: If the tensor corresponding to the input
-                :class:`TensorAttr` was not found.
         """
         attr = self._tensor_attr_cls.cast(*args, **kwargs)
         if not attr.is_fully_specified():
@@ -339,9 +330,9 @@ class FeatureStore(ABC):
                              f"specifying all 'UNSET' fields.")
 
         tensor = self._get_tensor(attr)
-        if tensor is None:
-            raise KeyError(f"A tensor corresponding to '{attr}' was not found")
-        return self._to_type(attr, tensor) if convert_type else tensor
+        if convert_type:
+            tensor = self._to_type(attr, tensor)
+        return tensor
 
     def _multi_get_tensor(
         self,
@@ -375,8 +366,6 @@ class FeatureStore(ABC):
         Raises:
             ValueError: If any input :class:`TensorAttr` is not fully
                 specified.
-            KeyError: If any of the tensors corresponding to the input
-                :class:`TensorAttr` was not found.
         """
         attrs = [self._tensor_attr_cls.cast(attr) for attr in attrs]
         bad_attrs = [attr for attr in attrs if not attr.is_fully_specified()]
@@ -387,20 +376,16 @@ class FeatureStore(ABC):
                 f"'UNSET' fields")
 
         tensors = self._multi_get_tensor(attrs)
-        if any(v is None for v in tensors):
-            bad_attrs = [attrs[i] for i, v in enumerate(tensors) if v is None]
-            raise KeyError(f"Tensors corresponding to attributes "
-                           f"'{bad_attrs}' were not found")
-
-        return [
-            self._to_type(attr, tensor) if convert_type else tensor
-            for attr, tensor in zip(attrs, tensors)
-        ]
+        if convert_type:
+            tensors = [
+                self._to_type(attr, tensor)
+                for attr, tensor in zip(attrs, tensors)
+            ]
+        return tensors
 
     @abstractmethod
     def _remove_tensor(self, attr: TensorAttr) -> bool:
         r"""To be implemented by :obj:`FeatureStore` subclasses."""
-        pass
 
     def remove_tensor(self, *args, **kwargs) -> bool:
         r"""Removes a tensor from the :class:`FeatureStore`.
@@ -458,7 +443,6 @@ class FeatureStore(ABC):
     @abstractmethod
     def get_all_tensor_attrs(self) -> List[TensorAttr]:
         r"""Returns all registered tensor attributes."""
-        pass
 
     # `AttrView` methods ######################################################
 
@@ -476,11 +460,9 @@ class FeatureStore(ABC):
         attr: TensorAttr,
         tensor: FeatureTensorType,
     ) -> FeatureTensorType:
-        if (isinstance(attr.index, torch.Tensor)
-                and isinstance(tensor, np.ndarray)):
+        if isinstance(attr.index, Tensor) and isinstance(tensor, np.ndarray):
             return torch.from_numpy(tensor)
-        if (isinstance(attr.index, np.ndarray)
-                and isinstance(tensor, torch.Tensor)):
+        if isinstance(attr.index, np.ndarray) and isinstance(tensor, Tensor):
             return tensor.detach().cpu().numpy()
         return tensor
 
@@ -491,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:
@@ -515,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

@@ -25,10 +25,10 @@ from typing import Any, Dict, List, Optional, Tuple
 
 from torch import Tensor
 
+from torch_geometric.index import index2ptr, ptr2index
 from torch_geometric.typing import EdgeTensorType, EdgeType, OptTensor
 from torch_geometric.utils import index_sort
 from torch_geometric.utils.mixin import CastMixin
-from torch_geometric.utils.sparse import index2ptr, ptr2index
 
 # The output of converting between two types in the GraphStore is a Tuple of
 # dictionaries: row, col, and perm. The dictionaries are keyed by the edge
@@ -116,7 +116,6 @@ class GraphStore(ABC):
     def _put_edge_index(self, edge_index: EdgeTensorType,
                         edge_attr: EdgeAttr) -> bool:
         r"""To be implemented by :class:`GraphStore` subclasses."""
-        pass
 
     def put_edge_index(self, edge_index: EdgeTensorType, *args,
                        **kwargs) -> bool:
@@ -137,7 +136,6 @@ class GraphStore(ABC):
     @abstractmethod
     def _get_edge_index(self, edge_attr: EdgeAttr) -> Optional[EdgeTensorType]:
         r"""To be implemented by :class:`GraphStore` subclasses."""
-        pass
 
     def get_edge_index(self, *args, **kwargs) -> EdgeTensorType:
         r"""Synchronously obtains an :obj:`edge_index` tuple from the
@@ -160,7 +158,6 @@ class GraphStore(ABC):
     @abstractmethod
     def _remove_edge_index(self, edge_attr: EdgeAttr) -> bool:
         r"""To be implemented by :class:`GraphStore` subclasses."""
-        pass
 
     def remove_edge_index(self, *args, **kwargs) -> bool:
         r"""Synchronously deletes an :obj:`edge_index` tuple from the
@@ -177,7 +174,6 @@ class GraphStore(ABC):
     @abstractmethod
     def get_all_edge_attrs(self) -> List[EdgeAttr]:
         r"""Returns all registered edge attributes."""
-        pass
 
     # Layout Conversion #######################################################
 

torch_geometric/data/hetero_data.py

@@ -10,6 +10,7 @@ import torch
 from torch import Tensor
 from typing_extensions import Self
 
+from torch_geometric import Index
 from torch_geometric.data import EdgeAttr, FeatureStore, GraphStore, TensorAttr
 from torch_geometric.data.data import BaseData, Data, size_repr, warn_or_raise
 from torch_geometric.data.graph_store import EdgeLayout
@@ -36,6 +37,8 @@ from torch_geometric.utils import (
 
 NodeOrEdgeStorage = Union[NodeStorage, EdgeStorage]
 
+_DISPLAYED_TYPE_NAME_WARNING: bool = False
+
 
 class HeteroData(BaseData, FeatureStore, GraphStore):
     r"""A data object describing a heterogeneous graph, holding multiple node
@@ -334,7 +337,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
     def __cat_dim__(self, key: str, value: Any,
                     store: Optional[NodeOrEdgeStorage] = None, *args,
                     **kwargs) -> Any:
-        if is_sparse(value) and 'adj' in key:
+        if is_sparse(value) and ('adj' in key or 'edge_index' in key):
             return (0, 1)
         elif isinstance(store, EdgeStorage) and 'index' in key:
             return -1
@@ -344,6 +347,8 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
                 store: Optional[NodeOrEdgeStorage] = None, *args,
                 **kwargs) -> Any:
         if 'batch' in key and isinstance(value, Tensor):
+            if isinstance(value, Index):
+                return value.get_dim_size()
             return int(value.max()) + 1
         elif isinstance(store, EdgeStorage) and 'index' in key:
             return torch.tensor(store.size()).view(2, 1)
@@ -562,11 +567,15 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         return mapping
 
     def _check_type_name(self, name: str):
-        if '__' in name:
-            warnings.warn(f"The type '{name}' contains double underscores "
-                          f"('__') which may lead to unexpected behavior. "
-                          f"To avoid any issues, ensure that your type names "
-                          f"only contain single underscores.")
+        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.")
 
     def get_node_store(self, key: NodeType) -> NodeStorage:
         r"""Gets the :class:`~torch_geometric.data.storage.NodeStorage` object
@@ -771,8 +780,8 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         for edge_type in self.edge_types:
             if edge_type not in edge_types:
                 del data[edge_type]
-        node_types = set(e[0] for e in edge_types)
-        node_types |= set(e[-1] for e in edge_types)
+        node_types = {e[0] for e in edge_types}
+        node_types |= {e[-1] for e in edge_types}
         for node_type in self.node_types:
             if node_type not in node_types:
                 del data[node_type]
@@ -878,7 +887,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
                 if len(sizes) != len(stores):
                     continue
                 # The attributes needs to have the same number of dimensions:
-                lengths = set([len(size) for size in sizes])
+                lengths = {len(size) for size in sizes}
                 if len(lengths) != 1:
                     continue
                 # The attributes needs to have the same size in all dimensions:

torch_geometric/data/in_memory_dataset.py

@@ -347,10 +347,8 @@ class InMemoryDataset(Dataset):
 def nested_iter(node: Union[Mapping, Sequence]) -> Iterable:
     if isinstance(node, Mapping):
         for key, value in node.items():
-            for inner_key, inner_value in nested_iter(value):
-                yield inner_key, inner_value
+            yield from nested_iter(value)
     elif isinstance(node, Sequence):
-        for i, inner_value in enumerate(node):
-            yield i, inner_value
+        yield from enumerate(node)
     else:
         yield None, node