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

{pyg_nightly-2.6.0.dev20240704.dist-info → pyg_nightly-2.8.0.dev20251207.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: flit 3.9.0
+Generator: flit 3.12.0
 Root-Is-Purelib: true
 Tag: py3-none-any

pyg_nightly-2.8.0.dev20251207.dist-info/licenses/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2023 PyG Team <team@pyg.org>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

torch_geometric/__init__.py

@@ -1,6 +1,13 @@
+from collections import defaultdict
+
+import torch
+import torch_geometric.typing
+
 from ._compile import compile, is_compiling
+from ._onnx import is_in_onnx_export, safe_onnx_export
 from .index import Index
 from .edge_index import EdgeIndex
+from .hash_tensor import HashTensor
 from .seed import seed_everything
 from .home import get_home_dir, set_home_dir
 from .device import is_mps_available, is_xpu_available, device
@@ -24,16 +31,19 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.6.0.dev20240704'
+__version__ = '2.8.0.dev20251207'
 
 __all__ = [
     'Index',
     'EdgeIndex',
+    'HashTensor',
     'seed_everything',
     'get_home_dir',
     'set_home_dir',
     'compile',
     'is_compiling',
+    'is_in_onnx_export',
+    'safe_onnx_export',
     'is_mps_available',
     'is_xpu_available',
     'device',
@@ -47,3 +57,26 @@ __all__ = [
     'torch_geometric',
     '__version__',
 ]
+
+if not torch_geometric.typing.WITH_PT113:
+    import warnings as std_warnings
+
+    std_warnings.warn(
+        "PyG 2.7 removed support for PyTorch < 1.13. Consider "
+        "Consider upgrading to PyTorch >= 1.13 or downgrading "
+        "to PyG <= 2.6. ", stacklevel=2)
+
+# Serialization ###############################################################
+
+if torch_geometric.typing.WITH_PT24:
+    torch.serialization.add_safe_globals([
+        dict,
+        list,
+        defaultdict,
+        Index,
+        torch_geometric.index.CatMetadata,
+        EdgeIndex,
+        torch_geometric.edge_index.SortOrder,
+        torch_geometric.edge_index.CatMetadata,
+        HashTensor,
+    ])

torch_geometric/_compile.py

@@ -10,6 +10,8 @@ def is_compiling() -> bool:
     r"""Returns :obj:`True` in case :pytorch:`PyTorch` is compiling via
     :meth:`torch.compile`.
     """
+    if torch_geometric.typing.WITH_PT23:
+        return torch.compiler.is_compiling()
     if torch_geometric.typing.WITH_PT21:
         return torch._dynamo.is_compiling()
     return False  # pragma: no cover
@@ -25,10 +27,16 @@ def compile(
     This function has the same signature as :meth:`torch.compile` (see
     `here <https://pytorch.org/docs/stable/generated/torch.compile.html>`__).
 
+    Args:
+        model: The model to compile.
+        *args: Additional arguments of :meth:`torch.compile`.
+        **kwargs: Additional keyword arguments of :meth:`torch.compile`.
+
     .. note::
         :meth:`torch_geometric.compile` is deprecated in favor of
         :meth:`torch.compile`.
     """
-    warnings.warn("'torch_geometric.compile' is deprecated in favor of "
-                  "'torch.compile'")
-    return torch.compile(model, *args, **kwargs)
+    warnings.warn(
+        "'torch_geometric.compile' is deprecated in favor of "
+        "'torch.compile'", stacklevel=2)
+    return torch.compile(model, *args, **kwargs)  # type: ignore

torch_geometric/_onnx.py

@@ -0,0 +1,228 @@
+import warnings
+from os import PathLike
+from typing import Any, Union
+
+import torch
+
+from torch_geometric import is_compiling
+
+
+def is_in_onnx_export() -> bool:
+    r"""Returns :obj:`True` in case :pytorch:`PyTorch` is exporting to ONNX via
+    :meth:`torch.onnx.export`.
+    """
+    if is_compiling():
+        return False
+    if torch.jit.is_scripting():
+        return False
+    return torch.onnx.is_in_onnx_export()
+
+
+def safe_onnx_export(
+    model: torch.nn.Module,
+    args: Union[torch.Tensor, tuple[Any, ...]],
+    f: Union[str, PathLike[Any], None],
+    skip_on_error: bool = False,
+    **kwargs: Any,
+) -> bool:
+    r"""A safe wrapper around :meth:`torch.onnx.export` that handles known
+    ONNX serialization issues in PyTorch Geometric.
+
+    This function provides workarounds for the ``onnx_ir.serde.SerdeError``
+    with boolean ``allowzero`` attributes that occurs in certain environments.
+
+    Args:
+        model (torch.nn.Module): The model to export.
+        args (torch.Tensor or tuple): The input arguments for the model.
+        f (str or PathLike): The file path to save the model.
+        skip_on_error (bool): If True, return False instead of raising when
+            workarounds fail. Useful for CI environments.
+        **kwargs: Additional arguments passed to :meth:`torch.onnx.export`.
+
+    Returns:
+        bool: True if export succeeded, False if skipped due to known issues
+              (only when skip_on_error=True).
+
+    Example:
+        >>> from torch_geometric.nn import SAGEConv
+        >>> from torch_geometric import safe_onnx_export
+        >>>
+        >>> class MyModel(torch.nn.Module):
+        ...     def __init__(self):
+        ...         super().__init__()
+        ...         self.conv = SAGEConv(8, 16)
+        ...     def forward(self, x, edge_index):
+        ...         return self.conv(x, edge_index)
+        >>>
+        >>> model = MyModel()
+        >>> x = torch.randn(3, 8)
+        >>> edge_index = torch.tensor([[0, 1, 2], [1, 0, 2]])
+        >>> success = safe_onnx_export(model, (x, edge_index), 'model.onnx')
+        >>>
+        >>> # For CI environments:
+        >>> success = safe_onnx_export(model, (x, edge_index), 'model.onnx',
+        ...                             skip_on_error=True)
+        >>> if not success:
+        ...     print("ONNX export skipped due to known upstream issue")
+    """
+    # Convert single tensor to tuple for torch.onnx.export compatibility
+    if isinstance(args, torch.Tensor):
+        args = (args, )
+
+    try:
+        # First attempt: standard ONNX export
+        torch.onnx.export(model, args, f, **kwargs)
+        return True
+
+    except Exception as e:
+        error_str = str(e)
+        error_type = type(e).__name__
+
+        # Check for the specific onnx_ir.serde.SerdeError patterns
+        is_allowzero_error = (('onnx_ir.serde.SerdeError' in error_str
+                               and 'allowzero' in error_str) or
+                              'ValueError: Value out of range: 1' in error_str
+                              or 'serialize_model_into' in error_str
+                              or 'serialize_attribute_into' in error_str)
+
+        if is_allowzero_error:
+            warnings.warn(
+                f"Encountered known ONNX serialization issue ({error_type}). "
+                "This is likely the allowzero boolean attribute bug. "
+                "Attempting workaround...", UserWarning, stacklevel=2)
+
+            # Apply workaround strategies
+            return _apply_onnx_allowzero_workaround(model, args, f,
+                                                    skip_on_error, **kwargs)
+
+        else:
+            # Re-raise other errors
+            raise
+
+
+def _apply_onnx_allowzero_workaround(
+    model: torch.nn.Module,
+    args: tuple[Any, ...],
+    f: Union[str, PathLike[Any], None],
+    skip_on_error: bool = False,
+    **kwargs: Any,
+) -> bool:
+    r"""Apply workaround strategies for onnx_ir.serde.SerdeError with allowzero
+    attributes.
+
+    Returns:
+        bool: True if export succeeded, False if skipped (when
+              skip_on_error=True).
+    """
+    # Strategy 1: Try without dynamo if it was enabled
+    if kwargs.get('dynamo', False):
+        try:
+            kwargs_no_dynamo = kwargs.copy()
+            kwargs_no_dynamo['dynamo'] = False
+
+            warnings.warn(
+                "Retrying ONNX export with dynamo=False as workaround",
+                UserWarning, stacklevel=3)
+
+            torch.onnx.export(model, args, f, **kwargs_no_dynamo)
+            return True
+
+        except Exception:
+            pass
+
+    # Strategy 2: Try with different opset versions
+    original_opset = kwargs.get('opset_version', 18)
+    for opset_version in [17, 16, 15, 14, 13, 11]:
+        if opset_version != original_opset:
+            try:
+                kwargs_opset = kwargs.copy()
+                kwargs_opset['opset_version'] = opset_version
+
+                warnings.warn(
+                    f"Retrying ONNX export with opset_version={opset_version}",
+                    UserWarning, stacklevel=3)
+
+                torch.onnx.export(model, args, f, **kwargs_opset)
+                return True
+
+            except Exception:
+                continue
+
+    # Strategy 3: Try legacy export (non-dynamo with older opset)
+    try:
+        kwargs_legacy = kwargs.copy()
+        kwargs_legacy['dynamo'] = False
+        kwargs_legacy['opset_version'] = 11
+
+        warnings.warn(
+            "Retrying ONNX export with legacy settings "
+            "(dynamo=False, opset_version=11)", UserWarning, stacklevel=3)
+
+        torch.onnx.export(model, args, f, **kwargs_legacy)
+        return True
+
+    except Exception:
+        pass
+
+    # Strategy 4: Try with minimal settings
+    try:
+        minimal_kwargs: dict[str, Any] = {
+            'opset_version': 11,
+            'dynamo': False,
+        }
+        # Add optional parameters if they exist
+        if kwargs.get('input_names') is not None:
+            minimal_kwargs['input_names'] = kwargs.get('input_names')
+        if kwargs.get('output_names') is not None:
+            minimal_kwargs['output_names'] = kwargs.get('output_names')
+
+        warnings.warn(
+            "Retrying ONNX export with minimal settings as last resort",
+            UserWarning, stacklevel=3)
+
+        torch.onnx.export(model, args, f, **minimal_kwargs)
+        return True
+
+    except Exception:
+        pass
+
+    # If all strategies fail, handle based on skip_on_error flag
+    import os
+    pytest_detected = 'PYTEST_CURRENT_TEST' in os.environ or 'pytest' in str(f)
+
+    if skip_on_error:
+        # For CI environments: skip gracefully instead of failing
+        warnings.warn(
+            "ONNX export skipped due to known upstream issue "
+            "(onnx_ir.serde.SerdeError). "
+            "This is caused by a bug in the onnx_ir package where boolean "
+            "allowzero attributes cannot be serialized. All workarounds "
+            "failed. Consider updating packages: pip install --upgrade onnx "
+            "onnxscript "
+            "onnx_ir", UserWarning, stacklevel=3)
+        return False
+
+    # For regular usage: provide detailed error message
+    error_msg = (
+        "Failed to export model to ONNX due to known serialization issue. "
+        "This is caused by a bug in the onnx_ir package where boolean "
+        "allowzero attributes cannot be serialized. "
+        "Workarounds attempted: dynamo=False, multiple opset versions, "
+        "and legacy export. ")
+
+    if pytest_detected:
+        error_msg += (
+            "\n\nThis error commonly occurs in pytest environments. "
+            "Try one of these solutions:\n"
+            "1. Run the export outside of pytest (in a regular Python "
+            "script)\n"
+            "2. Update packages: pip install --upgrade onnx onnxscript "
+            "onnx_ir\n"
+            "3. Use torch.jit.script() instead of ONNX export for testing\n"
+            "4. Use safe_onnx_export(..., skip_on_error=True) to skip "
+            "gracefully in CI")
+    else:
+        error_msg += ("\n\nTry updating packages: pip install --upgrade onnx "
+                      "onnxscript onnx_ir")
+
+    raise RuntimeError(error_msg)

torch_geometric/config_mixin.py

@@ -3,6 +3,8 @@ from dataclasses import fields, is_dataclass
 from importlib import import_module
 from typing import Any, Dict
 
+from torch.nn import ModuleDict, ModuleList
+
 from torch_geometric.config_store import (
     class_from_dataclass,
     dataclass_from_class,
@@ -71,9 +73,9 @@ def _recursive_config(value: Any) -> Any:
         return value.config()
     if is_torch_instance(value, ConfigMixin):
         return value.config()
-    if isinstance(value, (tuple, list)):
+    if isinstance(value, (tuple, list, ModuleList)):
         return [_recursive_config(v) for v in value]
-    if isinstance(value, dict):
+    if isinstance(value, (dict, ModuleDict)):
         return {k: _recursive_config(v) for k, v in value.items()}
     return value
 
@@ -82,7 +84,10 @@ def _recursive_from_config(value: Any) -> Any:
     cls: Any = None
     if is_dataclass(value):
         if getattr(value, '_target_', None):
-            cls = _locate_cls(value._target_)
+            try:
+                cls = _locate_cls(value._target_)  # type: ignore
+            except ImportError:
+                pass  # Keep the dataclass as it is.
         else:
             cls = class_from_dataclass(value.__class__)
     elif isinstance(value, dict) and '_target_' in value:

torch_geometric/config_store.py

@@ -168,7 +168,7 @@ def map_annotation(
         assert origin is not None
         args = tuple(map_annotation(a, mapping) for a in args)
         if type(annotation).__name__ == 'GenericAlias':
-            # If annotated with `list[...]` or `dict[...]` (>= Python 3.10):
+            # If annotated with `list[...]` or `dict[...]`:
             annotation = origin[args]
         else:
             # If annotated with `typing.List[...]` or `typing.Dict[...]`:

torch_geometric/contrib/__init__.py

@@ -7,6 +7,6 @@ import torch_geometric.contrib.explain  # noqa
 
 warnings.warn(
     "'torch_geometric.contrib' contains experimental code and is subject to "
-    "change. Please use with caution.")
+    "change. Please use with caution.", stacklevel=2)
 
 __all__ = []

torch_geometric/contrib/explain/pgm_explainer.py

@@ -151,7 +151,7 @@ class PGMExplainer(ExplainerAlgorithm):
 
             pred_change = torch.max(soft_pred) - soft_pred_perturb[pred_label]
 
-            sample[num_nodes] = pred_change
+            sample[num_nodes] = pred_change.detach()
             samples.append(sample)
 
         samples = torch.tensor(np.array(samples))

torch_geometric/data/__init__.py

@@ -1,7 +1,10 @@
 # flake8: noqa
 
+import torch
+import torch_geometric.typing
+
 from .feature_store import FeatureStore, TensorAttr
-from .graph_store import GraphStore, EdgeAttr
+from .graph_store import GraphStore, EdgeAttr, EdgeLayout
 from .data import Data
 from .hetero_data import HeteroData
 from .batch import Batch
@@ -68,6 +71,21 @@ from torch_geometric.loader import DataLoader
 from torch_geometric.loader import DataListLoader
 from torch_geometric.loader import DenseDataLoader
 
+# Serialization ###############################################################
+
+if torch_geometric.typing.WITH_PT24:
+    torch.serialization.add_safe_globals([
+        Data,
+        HeteroData,
+        TemporalData,
+        ClusterData,
+        TensorAttr,
+        EdgeAttr,
+        EdgeLayout,
+    ])
+
+# Deprecations ################################################################
+
 NeighborSampler = deprecated(  # type: ignore
     details="use 'loader.NeighborSampler' instead",
     func_name='data.NeighborSampler',

torch_geometric/data/batch.py

@@ -125,8 +125,8 @@ class Batch(metaclass=DynamicInheritance):
             cls=self.__class__.__bases__[-1],
             batch=self,
             idx=idx,
-            slice_dict=getattr(self, '_slice_dict'),
-            inc_dict=getattr(self, '_inc_dict'),
+            slice_dict=self._slice_dict,
+            inc_dict=self._inc_dict,
             decrement=True,
         )
 

torch_geometric/data/collate.py

@@ -191,10 +191,8 @@ def _collate(
             if torch_geometric.typing.WITH_PT20:
                 storage = elem.untyped_storage()._new_shared(
                     numel * elem.element_size(), device=elem.device)
-            elif torch_geometric.typing.WITH_PT112:
-                storage = elem.storage()._new_shared(numel, device=elem.device)
             else:
-                storage = elem.storage()._new_shared(numel)
+                storage = elem.storage()._new_shared(numel, device=elem.device)
             shape = list(elem.size())
             if cat_dim is None or elem.dim() == 0:
                 shape = [len(values)] + shape

torch_geometric/data/data.py

@@ -1,5 +1,6 @@
 import copy
 import warnings
+from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from itertools import chain
@@ -354,7 +355,7 @@ class BaseData:
         """
         return self.apply(lambda x: x.contiguous(), *args)
 
-    def to(self, device: Union[int, str], *args: str,
+    def to(self, device: Union[int, str, torch.device], *args: str,
            non_blocking: bool = False):
         r"""Performs tensor device conversion, either for all attributes or
         only the ones given in :obj:`*args`.
@@ -659,7 +660,13 @@ class Data(BaseData, FeatureStore, GraphStore):
                 return value.get_dim_size()
             return int(value.max()) + 1
         elif 'index' in key or key == 'face':
-            return self.num_nodes
+            num_nodes = self.num_nodes
+            if num_nodes is None:
+                raise RuntimeError(f"Unable to infer 'num_nodes' from the "
+                                   f"attribute '{key}'. Please explicitly set "
+                                   f"'num_nodes' as an attribute of 'data' to "
+                                   f"prevent this error")
+            return num_nodes
         else:
             return 0
 
@@ -844,14 +851,14 @@ class Data(BaseData, FeatureStore, GraphStore):
         # that maps global node indices to local ones in the final
         # heterogeneous graph:
         node_ids, index_map = {}, torch.empty_like(node_type)
-        for i, key in enumerate(node_type_names):
+        for i in range(len(node_type_names)):
             node_ids[i] = (node_type == i).nonzero(as_tuple=False).view(-1)
             index_map[node_ids[i]] = torch.arange(len(node_ids[i]),
                                                   device=index_map.device)
 
         # We iterate over edge types to find the local edge indices:
         edge_ids = {}
-        for i, key in enumerate(edge_type_names):
+        for i in range(len(edge_type_names)):
             edge_ids[i] = (edge_type == i).nonzero(as_tuple=False).view(-1)
 
         data = HeteroData()
@@ -898,6 +905,60 @@ class Data(BaseData, FeatureStore, GraphStore):
 
         return data
 
+    def connected_components(self) -> List[Self]:
+        r"""Extracts connected components of the graph using a union-find
+        algorithm. The components are returned as a list of
+        :class:`~torch_geometric.data.Data` objects, where each object
+        represents a connected component of the graph.
+
+        .. code-block::
+
+            data = Data()
+            data.x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
+            data.y = torch.tensor([[1.1], [2.1], [3.1], [4.1]])
+            data.edge_index = torch.tensor(
+                [[0, 1, 2, 3], [1, 0, 3, 2]], dtype=torch.long
+            )
+
+            components = data.connected_components()
+            print(len(components))
+            >>> 2
+
+            print(components[0].x)
+            >>> Data(x=[2, 1], y=[2, 1], edge_index=[2, 2])
+
+        Returns:
+            List[Data]: A list of disconnected components.
+        """
+        # Union-Find algorithm to find connected components
+        self._parents: Dict[int, int] = {}
+        self._ranks: Dict[int, int] = {}
+        for edge in self.edge_index.t().tolist():
+            self._union(edge[0], edge[1])
+
+        # Rerun _find_parent to ensure all nodes are covered correctly
+        for node in range(self.num_nodes):
+            self._find_parent(node)
+
+        # Group parents
+        grouped_parents = defaultdict(list)
+        for node, parent in self._parents.items():
+            grouped_parents[parent].append(node)
+        del self._ranks
+        del self._parents
+
+        # Create components based on the found parents (roots)
+        components: List[Self] = []
+        for nodes in grouped_parents.values():
+            # Convert the list of node IDs to a tensor
+            subset = torch.tensor(nodes, dtype=torch.long)
+
+            # Use the existing subgraph function
+            component_data = self.subgraph(subset)
+            components.append(component_data)
+
+        return components
+
     ###########################################################################
 
     @classmethod
@@ -1144,6 +1205,49 @@ class Data(BaseData, FeatureStore, GraphStore):
 
         return list(edge_attrs.values())
 
+    # Connected Components Helper Functions ###################################
+
+    def _find_parent(self, node: int) -> 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 (int): The node for which to find the representative parent.
+
+        Returns:
+            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: int, node2: int):
+        r"""Merges the sets containing 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 (int): The index of the first node to union.
+            node2 (int): The index of the second node to union.
+        """
+        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
+
 
 ###############################################################################
 
@@ -1165,7 +1269,7 @@ def size_repr(key: Any, value: Any, indent: int = 0) -> str:
                f'[{value.num_rows}, {value.num_cols}])')
     elif isinstance(value, str):
         out = f"'{value}'"
-    elif isinstance(value, Sequence):
+    elif isinstance(value, (Sequence, set)):
         out = str([len(value)])
     elif isinstance(value, Mapping) and len(value) == 0:
         out = '{}'
@@ -1187,4 +1291,4 @@ def warn_or_raise(msg: str, raise_on_error: bool = True):
     if raise_on_error:
         raise ValueError(msg)
     else:
-        warnings.warn(msg)
+        warnings.warn(msg, stacklevel=2)