pyg-nightly 2.7.0.dev20250903__py3-none-any.whl → 2.7.0.dev20250905__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250903.dist-info → pyg_nightly-2.7.0.dev20250905.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250903
+Version: 2.7.0.dev20250905
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>

{pyg_nightly-2.7.0.dev20250903.dist-info → pyg_nightly-2.7.0.dev20250905.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=FdFwyBPcM-xmx9bjJNEjeOhn8oy2Hji4JY1KU0OrrH4,2292
+torch_geometric/__init__.py,sha256=JaI2udwlYGKK9_OGrMEQn05CBfOSEPFFsk9OtBnPN2c,2292
 torch_geometric/_compile.py,sha256=9yqMTBKatZPr40WavJz9FjNi7pQj8YZAZOyZmmRGXgc,1351
 torch_geometric/_onnx.py,sha256=ODB_8cwFUiwBUjngXn6-K5HHb7IDul7DDXuuGX7vj_0,8178
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -33,7 +33,7 @@ torch_geometric/contrib/transforms/__init__.py,sha256=lrGnWsEiJf5zsBRmshGZZFN_uY
 torch_geometric/data/__init__.py,sha256=D6Iz5A9vEb_2rpf96Zn7uM-lchZ3WpW8X7WdAD1yxKw,4565
 torch_geometric/data/batch.py,sha256=8X8CN4_1rjrh48R3R2--mZUgfsO7Po9JP-H6SbrBiBA,8740
 torch_geometric/data/collate.py,sha256=tOUvttXoEo-bOvJx_qMivJq2JqOsB9iDdjovtiyys4o,12644
-torch_geometric/data/data.py,sha256=F6DLTd-hJI9hiHemS04Eh-kGZTbg_16-gvRVFsbXfRM,43827
+torch_geometric/data/data.py,sha256=-E6el1knNgSJyapV8KUk2aRRHOfvwEvjUFfe_BapLfc,47490
 torch_geometric/data/database.py,sha256=K3KLefYVfsBN9HRItgFZNkbUIllfDt4ueauBFxk3Rxk,23106
 torch_geometric/data/datapipes.py,sha256=9_Cq3j_7LIF4plQFzbLaqyy0LcpKdAic6yiKgMqSX9A,3083
 torch_geometric/data/dataset.py,sha256=AaJH0N9eZgvxX0ljyTH8cXutKJ0AGFAyE-H4Sw9D51w,16834
@@ -41,7 +41,7 @@ torch_geometric/data/download.py,sha256=kcesTu6jlgmCeePpOxDQOnVhxB_GuZ9iu9ds72KE
 torch_geometric/data/extract.py,sha256=DMG8_6ps4O5xKfkb7j1gUBX_jlWpFdmz6OLY2jBSEx4,2339
 torch_geometric/data/feature_store.py,sha256=pl2pJL25wqzEZnNZbW8c8Ee_yH0DnE2AK8TioTWZV-g,20045
 torch_geometric/data/graph_store.py,sha256=dSMCcMYlka2elfw-Rof-lG_iGQv6NHX98uPEVcgDn_g,13900
-torch_geometric/data/hetero_data.py,sha256=2LV8pSvv-IWkTUk8xg7VeI17YMhikg1RkeQhMwA8tkE,48583
+torch_geometric/data/hetero_data.py,sha256=i4_Wt8MFyEG1ZhsU_zg7IKqA2ZatrYyjtYs_R_Qrf3U,53353
 torch_geometric/data/hypergraph_data.py,sha256=LfriiuJRx9ZrrSGj_fO5NUsh4kvyXJuRdCOqsWo__vc,8307
 torch_geometric/data/in_memory_dataset.py,sha256=ilFxjF4pvBILsS4wOqocwRBc2j6toI2S_KMFF19KB1w,13413
 torch_geometric/data/large_graph_indexer.py,sha256=myXTXhbRHQPxEOHNHPeNHB_pBzXCIQBr1KQt9WwBoi8,25468
@@ -294,7 +294,7 @@ torch_geometric/loader/temporal_dataloader.py,sha256=Z7L_rYdl6SYBQXAgtr18FVcmfMH
 torch_geometric/loader/utils.py,sha256=3hzKzIgB52QIZu7Jdn4JeXZaegIJinIQfIUP9DrUWUQ,14903
 torch_geometric/loader/zip_loader.py,sha256=3lt10fD15Rxm1WhWzypswGzCEwUz4h8OLCD1nE15yNg,3843
 torch_geometric/metrics/__init__.py,sha256=3krvDobW6vV5yHTjq2S2pmOXxNfysNG26muq7z48e94,699
-torch_geometric/metrics/link_pred.py,sha256=mRQTSYYJgLKXFCelZHMKVOSbPED11JVhbryp7ajjxDU,31137
+torch_geometric/metrics/link_pred.py,sha256=1_hE3KiRqAdZLI6QuUbjgyFC__mTyFu_RimM3bD8wRw,31678
 torch_geometric/nn/__init__.py,sha256=kQHHHUxFDht2ztD-XFQuv98TvC8MdodaFsIjAvltJBw,874
 torch_geometric/nn/data_parallel.py,sha256=YiybTWoSFyfSzlXAamZ_-y1f7B6tvDEFHOuy_AyJz9Q,4761
 torch_geometric/nn/encoding.py,sha256=3DCOCO-XFt-lMb97sHWGN-4KeGUFY5lVo9P00SzrCNk,3559
@@ -620,7 +620,7 @@ torch_geometric/utils/_trim_to_layer.py,sha256=cauOEzMJJK4w9BC-Pg1bHVncBYqG9XxQe
 torch_geometric/utils/_unbatch.py,sha256=B0vjKI96PtHvSBG8F_lqvsiJE134aVjUurPZsG6UZRI,2378
 torch_geometric/utils/augmentation.py,sha256=1F0YCuaklZ9ZbXxdFV0oOoemWvLd8p60WvFo2chzl7E,8600
 torch_geometric/utils/convert.py,sha256=RE5n5no74Xu39-QMWFE0-1RvTgykdK33ymyjF9WcuSs,21938
-torch_geometric/utils/cross_entropy.py,sha256=ZFS5bivtzv3EV9zqgKsekmuQyoZZggPSclhl_tRNHxo,3047
+torch_geometric/utils/cross_entropy.py,sha256=_6whuSCWKNzavOLf83uZbI9RFU0wGRHQict_4-8XIYs,3773
 torch_geometric/utils/dropout.py,sha256=gg0rDnD4FLvBaKSoLAkZwViAQflhLefJm6_Mju5dmQs,11416
 torch_geometric/utils/embedding.py,sha256=Ac_MPSrZGpw-e-gU6Yz-seUioC2WZxBSSzXFeclGwMk,5232
 torch_geometric/utils/functions.py,sha256=orQdS_6EpzWSmBHSok3WhxCzLy9neB-cin1aTnlXY-8,703
@@ -646,7 +646,7 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=b-HnVesXjyJ_L1N-DnjiRiRVf7lhwKaBQF_2i5YMVSU,208
 torch_geometric/visualization/graph.py,sha256=mfZHXYfiU-CWMtfawYc80IxVwVmtK9hbIkSKhM_j7oI,14311
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250903.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250903.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250903.dist-info/METADATA,sha256=vYdkEEMQ1rrRaIksgih36EJ9htFz8pD4ZkQM7Zb3UsA,64100
-pyg_nightly-2.7.0.dev20250903.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250905.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250905.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250905.dist-info/METADATA,sha256=7zFs8SgMgRdkwK679QLO1W08UGscHMakHb-3NXsiEok,64100
+pyg_nightly-2.7.0.dev20250905.dist-info/RECORD,,

torch_geometric/__init__.py

@@ -31,7 +31,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250903'
+__version__ = '2.7.0.dev20250905'
 
 __all__ = [
     'Index',

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
@@ -904,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
@@ -1150,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
+
 
 ###############################################################################
 

torch_geometric/data/hetero_data.py

@@ -487,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
 
@@ -1148,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/metrics/link_pred.py

@@ -21,6 +21,19 @@ class LinkPredMetricData:
     edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]]
     edge_label_weight: Optional[Tensor] = None
 
+    def __post_init__(self) -> None:
+        # Filter all negative weights - they should not be used as ground-truth
+        if self.edge_label_weight is not None:
+            pos_mask = self.edge_label_weight > 0
+            self.edge_label_weight = self.edge_label_weight[pos_mask]
+            if isinstance(self.edge_label_index, Tensor):
+                self.edge_label_index = self.edge_label_index[:, pos_mask]
+            else:
+                self.edge_label_index = (
+                    self.edge_label_index[0][pos_mask],
+                    self.edge_label_index[1][pos_mask],
+                )
+
     @property
     def pred_rel_mat(self) -> Tensor:
         r"""Returns a matrix indicating the relevance of the `k`-th prediction.
@@ -374,8 +387,6 @@ class LinkPredMetricCollection(torch.nn.ModuleDict):
         if self.weighted and edge_label_weight is None:
             raise ValueError(f"'edge_label_weight' is a required argument for "
                              f"weighted '{self.__class__.__name__}' metrics")
-        if not self.weighted:
-            edge_label_weight = None
 
         data = LinkPredMetricData(  # Share metric data across metrics.
             pred_index_mat=pred_index_mat,

torch_geometric/utils/cross_entropy.py

@@ -18,30 +18,51 @@ class SparseCrossEntropy(torch.autograd.Function):
     ) -> Tensor:
         assert inputs.dim() == 2
 
-        logsumexp = inputs.logsumexp(dim=-1)
-        ctx.save_for_backward(inputs, edge_label_index, edge_label_weight,
-                              logsumexp)
+        # Support for both positive and negative weights:
+        # Positive weights scale the logits *after* softmax.
+        # Negative weights scale the denominator *before* softmax:
+        pos_y = edge_label_index
+        neg_y = pos_weight = neg_weight = None
 
-        out = inputs[edge_label_index[0], edge_label_index[1]]
-        out.neg_().add_(logsumexp[edge_label_index[0]])
         if edge_label_weight is not None:
-            out *= edge_label_weight
+            pos_mask = edge_label_weight >= 0
+            pos_y = edge_label_index[:, pos_mask]
+            pos_weight = edge_label_weight[pos_mask]
+
+            if pos_y.size(1) < edge_label_index.size(1):
+                neg_mask = ~pos_mask
+                neg_y = edge_label_index[:, neg_mask]
+                neg_weight = edge_label_weight[neg_mask]
+
+            if neg_y is not None and neg_weight is not None:
+                inputs = inputs.clone()
+                inputs[
+                    neg_y[0],
+                    neg_y[1],
+                ] += neg_weight.abs().log().clamp(min=1e-12)
+
+        logsumexp = inputs.logsumexp(dim=-1)
+        ctx.save_for_backward(inputs, pos_y, pos_weight, logsumexp)
+
+        out = inputs[pos_y[0], pos_y[1]]
+        out.neg_().add_(logsumexp[pos_y[0]])
+        if pos_weight is not None:
+            out *= pos_weight
 
         return out.sum() / inputs.size(0)
 
     @staticmethod
     @torch.autograd.function.once_differentiable
     def backward(ctx: Any, grad_out: Tensor) -> Tuple[Tensor, None, None]:
-        inputs, edge_label_index, edge_label_weight, logsumexp = (
-            ctx.saved_tensors)
+        inputs, pos_y, pos_weight, logsumexp = ctx.saved_tensors
 
         grad_out = grad_out / inputs.size(0)
-        grad_out = grad_out.expand(edge_label_index.size(1))
+        grad_out = grad_out.expand(pos_y.size(1))
 
-        if edge_label_weight is not None:
-            grad_out = grad_out * edge_label_weight
+        if pos_weight is not None:
+            grad_out = grad_out * pos_weight
 
-        grad_logsumexp = scatter(grad_out, edge_label_index[0], dim=0,
+        grad_logsumexp = scatter(grad_out, pos_y[0], dim=0,
                                  dim_size=inputs.size(0), reduce='sum')
 
         # Gradient computation of `logsumexp`: `grad * (self - result).exp()`
@@ -49,7 +70,7 @@ class SparseCrossEntropy(torch.autograd.Function):
         grad_input.exp_()
         grad_input.mul_(grad_logsumexp.view(-1, 1))
 
-        grad_input[edge_label_index[0], edge_label_index[1]] -= grad_out
+        grad_input[pos_y[0], pos_y[1]] -= grad_out
 
         return grad_input, None, None