pyg-nightly 2.7.0.dev20250902__py3-none-any.whl → 2.7.0.dev20250904__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250902.dist-info → pyg_nightly-2.7.0.dev20250904.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250902
+Version: 2.7.0.dev20250904
 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.dev20250902.dist-info → pyg_nightly-2.7.0.dev20250904.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=NO3jwSl5Hf33qP9GchFcqCCeokVu342RzUIGxdqipDI,2292
+torch_geometric/__init__.py,sha256=kAQ6fnK2P5phSQX95Xx4qdMo6r6k-i2p8cDLRERlxuw,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,15 +33,15 @@ 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
 torch_geometric/data/download.py,sha256=kcesTu6jlgmCeePpOxDQOnVhxB_GuZ9iu9ds72KEORc,1889
 torch_geometric/data/extract.py,sha256=DMG8_6ps4O5xKfkb7j1gUBX_jlWpFdmz6OLY2jBSEx4,2339
-torch_geometric/data/feature_store.py,sha256=BIMgIWpP1y7OCIQxnkdSWcnm8_BFJXuS_zOqfZZQOjI,20045
-torch_geometric/data/graph_store.py,sha256=EtIgsyY7RdBHRTCn34VypEBOG8cg8WzsNT_kTFKxJR4,13900
-torch_geometric/data/hetero_data.py,sha256=2LV8pSvv-IWkTUk8xg7VeI17YMhikg1RkeQhMwA8tkE,48583
+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=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
@@ -57,7 +57,7 @@ torch_geometric/data/lightning/__init__.py,sha256=w3En1tJfy3kSqe1MycpOyZpHFO3fxB
 torch_geometric/data/lightning/datamodule.py,sha256=jDv9ibLQV_FyPeq8ncq77oOU8qy-STCf-aaYd0R8JE8,29545
 torch_geometric/datasets/__init__.py,sha256=rgfUmjd9U3o8renKVl81Brscx4LOtwWmt6qAoaG41C4,6417
 torch_geometric/datasets/actor.py,sha256=oUxgJIX8bi5hJr1etWNYIFyVQNDDXi1nyVpHGGMEAGQ,4304
-torch_geometric/datasets/airfrans.py,sha256=8cCBmHPttrlKY_iwfyr-K-CUX_JEDjrIOg3r9dQSN7o,5439
+torch_geometric/datasets/airfrans.py,sha256=Pc9C7IuEKkKzko_RmFPQ5gzOAGJ3132DoZZ4HaePBT8,5440
 torch_geometric/datasets/airports.py,sha256=b3gkv3gY2JkUpmGiz36Z-g7EcnSfU8lBG1YsCOWdJ6k,3758
 torch_geometric/datasets/amazon.py,sha256=zLiAgrd_44LAFb8drsrIphRJNyuWa6TMjZgmoWdf97Y,3005
 torch_geometric/datasets/amazon_book.py,sha256=I-8kRsKgk9M60D4icYDELajlsRwksMLDaHMyn6sBC1Y,3214
@@ -191,7 +191,7 @@ torch_geometric/distributed/dist_neighbor_sampler.py,sha256=YrL-NMFOJwHJpF189o4k
 torch_geometric/distributed/event_loop.py,sha256=wr3iwMYEWOGkBlvC5huD2k5YxisaGE9w1Z-8RcQiIQk,3309
 torch_geometric/distributed/local_feature_store.py,sha256=CLW37RN0ouDufEs2tY9d2nLLvpxubiT6zgW3LIHAU8k,19058
 torch_geometric/distributed/local_graph_store.py,sha256=wNHXSS824Kk2HynbtWFXx-W4pl97UUBv6qFHAVqO3W4,8445
-torch_geometric/distributed/partition.py,sha256=BgjmhDloaooAXM7onGizrcikZs8oRnz5drQZHPDDO_g,14734
+torch_geometric/distributed/partition.py,sha256=VYw_3CdpRKXr1O4C80JRSMm8Od6xrS3t6H2bfmfJlGE,14733
 torch_geometric/distributed/rpc.py,sha256=j4TZQkk7NB2CIovRrasyvL9l9L4J6_YOq43gpzFMxow,5713
 torch_geometric/distributed/utils.py,sha256=FGrr3qw7hx7EQaIjjqasurloCFJ9q_0jt8jdSIUjBeM,6567
 torch_geometric/explain/__init__.py,sha256=pRxVB33zsxhED1StRWdHboQWh3e06__g9N298Hzi42Y,359
@@ -271,7 +271,7 @@ torch_geometric/io/txt_array.py,sha256=LDeX2qtlNKW-kVe-wpnskMwAdXQp1jVCGQnrJce7S
 torch_geometric/loader/__init__.py,sha256=DJrdCD1A5PuBYRSgxFbZU9GTBStNuKngqkUV1oEfQQ4,1971
 torch_geometric/loader/base.py,sha256=ataIwNEYL0px3CN3LJEgXIVTRylDHB6-yBFXXuX2JN0,1615
 torch_geometric/loader/cache.py,sha256=S65heO3YTyUPbttqizCNtKPHIoAw5iHRpbvw6KlXmok,2106
-torch_geometric/loader/cluster.py,sha256=eMNxVkvZt5oQ_gJRgmWm1NBX7zU2tZI_BPaXeB0wuyk,13465
+torch_geometric/loader/cluster.py,sha256=CbZUy739vzMqOKgof2N73uc-Br4Daw56G3XMzptLUT8,13469
 torch_geometric/loader/data_list_loader.py,sha256=uLNqeMTkza8EEBjzqZWvgQS5kv5qWa9dyyxt6lIlcUA,1459
 torch_geometric/loader/dataloader.py,sha256=XzboK_Ygnzvaj2UQ1Q0az-6fdlKsUKlsbjo07sbErrQ,3527
 torch_geometric/loader/dense_data_loader.py,sha256=GDb_Vu2XyNL5iYzw2zoh1YiurZRr6d7mnT6HF2GWKxM,1685
@@ -377,7 +377,7 @@ torch_geometric/nn/conv/hgt_conv.py,sha256=lUhTWUMovMtn9yR_b2-kLNLqHChGOUl2OtXBY
 torch_geometric/nn/conv/hypergraph_conv.py,sha256=4BosbbqJyprlI6QjPqIfMxCqnARU_0mUn1zcAQhbw90,8691
 torch_geometric/nn/conv/le_conv.py,sha256=DonmmYZOKk5wIlTZzzIfNKqBY6MO0MRxYhyr0YtNz-Q,3494
 torch_geometric/nn/conv/lg_conv.py,sha256=8jMa79iPsOUbXEfBIc3wmbvAD8T3d1j37LeIFTX3Yag,2369
-torch_geometric/nn/conv/meshcnn_conv.py,sha256=92zUcgfS0Fwv-MpddF4Ia1a65y7ddPAkazYf7D6kvwg,21951
+torch_geometric/nn/conv/meshcnn_conv.py,sha256=qt9oAlj6krDU2DBkgr6s_dPw1_vtxfish4iW74JZ70g,21951
 torch_geometric/nn/conv/message_passing.py,sha256=ZuTvSvodGy1GyAW4mHtuoMUuxclam-7opidYNY5IHm8,44377
 torch_geometric/nn/conv/mf_conv.py,sha256=SkOGMN1tFT9dcqy8xYowsB2ozw6QfkoArgR1BksZZaU,4340
 torch_geometric/nn/conv/mixhop_conv.py,sha256=qVDPWeWcnO7_eHM0ZnpKtr8SISjb4jp0xjgpoDrwjlk,4555
@@ -463,11 +463,11 @@ torch_geometric/nn/models/molecule_gpt.py,sha256=k-XULH6jaurj-R2EE4sIWTkqlNqa3Cz
 torch_geometric/nn/models/neural_fingerprint.py,sha256=pTLJgU9Uh2Lnf9bggLj4cKI8YdEFcMF-9MALuubqbuQ,2378
 torch_geometric/nn/models/node2vec.py,sha256=81Ku4Rp4IwLEAy06KEgJ2fYtXXVL_uv_Hb8lBr6YXrE,7664
 torch_geometric/nn/models/pmlp.py,sha256=dcAASVSyQMMhItSfEJWPeAFh0R3tNCwAHwdrShwQ8o4,3538
-torch_geometric/nn/models/polynormer.py,sha256=mayWdzdolT5PCt_Oo7UGG-JUripMHWB2lUWF1bh6goU,7640
+torch_geometric/nn/models/polynormer.py,sha256=JgUngkF18sgepAAJTO7js9RISmYLWiO04-JEeV4J__8,7641
 torch_geometric/nn/models/protein_mpnn.py,sha256=SwTgafSbI2KJ-yqzn0trZtVWLmfo0_kPEaWSNJUCt70,12266
 torch_geometric/nn/models/re_net.py,sha256=pz66q5b5BoGDNVQvpEGS2RGoeKvpjkYAv9r3WAuvITk,8986
 torch_geometric/nn/models/rect.py,sha256=2F3XyyvHTAEuqfJpiNB5M8pSGy738LhPiom5I-SDWqM,2808
-torch_geometric/nn/models/rev_gnn.py,sha256=Bpme087Zs227lcB0ODOKWsxaly67q96wseaRt6bacjs,11796
+torch_geometric/nn/models/rev_gnn.py,sha256=bkKBAd_vXZ3UDMTJgdVObqleYHOTVsVcisftK7XoDlo,11797
 torch_geometric/nn/models/schnet.py,sha256=0aaHrVtxApdvn3RHCGLQJW1MbIb--CSYUrx9O3hDOZM,16656
 torch_geometric/nn/models/sgformer.py,sha256=3NDzkEVRtM1QmeJsXSq7FBhGGchyUvyX1SDPKYg9F70,6875
 torch_geometric/nn/models/signed_gcn.py,sha256=HEKaXZIWoDnsBRxIytviTpwsjQIFKl44c9glNUpwhlM,9841
@@ -626,7 +626,7 @@ torch_geometric/utils/embedding.py,sha256=Ac_MPSrZGpw-e-gU6Yz-seUioC2WZxBSSzXFec
 torch_geometric/utils/functions.py,sha256=orQdS_6EpzWSmBHSok3WhxCzLy9neB-cin1aTnlXY-8,703
 torch_geometric/utils/geodesic.py,sha256=e_XCn7dxqeYJBL-sAc2DfxF3kp_ZUIP0vwqsx1yshmU,4777
 torch_geometric/utils/hetero.py,sha256=ok4uAAOyMiaeEPmvyS4DNoDwdKnLS2gmgs5WVVklxOo,5539
-torch_geometric/utils/influence.py,sha256=7R-NW4myJMJPkbNiwcHTmO_m_B3gPB2IlBbQkB446xc,10348
+torch_geometric/utils/influence.py,sha256=wZqt4iR5s0t5LyB--srDthJnR8h9HQN3IqdVtjzD3Cc,10351
 torch_geometric/utils/isolated.py,sha256=nUxCfMY3q9IIFjelr4eyAJH4sYG9W3lGdpWidnp3dm4,3588
 torch_geometric/utils/laplacian.py,sha256=ludDil4yS1A27PEuYOjZtCtE3o-t0lnucJKfiqENhvM,3695
 torch_geometric/utils/loop.py,sha256=MUWUS7a5GxuxLKlCtRq95U1hc3MndybAhqKD5IAe2RY,23051
@@ -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.dev20250902.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250902.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250902.dist-info/METADATA,sha256=voU97tZI_J6eLyfyPVL0Q-7pwlS9ifAPhaMzuAsqiis,64100
-pyg_nightly-2.7.0.dev20250902.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250904.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250904.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250904.dist-info/METADATA,sha256=-vA9MGqxs1s0QEIIV_Vj3gZop1fABzvtq8fMNiso7jY,64100
+pyg_nightly-2.7.0.dev20250904.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.dev20250902'
+__version__ = '2.7.0.dev20250904'
 
 __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/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
@@ -352,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.
 
@@ -412,7 +412,7 @@ class FeatureStore(ABC):
         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:

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

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/datasets/airfrans.py

@@ -30,8 +30,8 @@ class AirfRANS(InMemoryDataset):
     divided by the specific mass (one component in meter squared per second
     squared), the turbulent kinematic viscosity (one component in meter squared
     per second).
-    Finaly, a boolean is attached to each point to inform if this point lies on
-    the airfoil or not.
+    Finally, a boolean is attached to each point to inform if this point lies
+    on the airfoil or not.
 
     A library for manipulating simulations of the dataset is available `here
     <https://airfrans.readthedocs.io/en/latest/index.html>`_.

torch_geometric/distributed/partition.py

@@ -361,7 +361,7 @@ class Partitioner:
             'edge_types': self.edge_types,
             'node_offset': list(node_offset.values()) if node_offset else None,
             'is_hetero': self.is_hetero,
-            'is_sorted': True,  # Based on columnn/destination.
+            'is_sorted': True,  # Based on column/destination.
         }
         with open(osp.join(self.root, 'META.json'), 'w') as f:
             json.dump(meta, f)

torch_geometric/loader/cluster.py

@@ -235,9 +235,9 @@ class ClusterData(torch.utils.data.Dataset):
 class ClusterLoader(torch.utils.data.DataLoader):
     r"""The data loader scheme from the `"Cluster-GCN: An Efficient Algorithm
     for Training Deep and Large Graph Convolutional Networks"
-    <https://arxiv.org/abs/1905.07953>`_ paper which merges partioned subgraphs
-    and their between-cluster links from a large-scale graph data object to
-    form a mini-batch.
+    <https://arxiv.org/abs/1905.07953>`_ paper which merges partitioned
+    subgraphs and their between-cluster links from a large-scale graph data
+    object to form a mini-batch.
 
     .. note::
 
@@ -252,7 +252,7 @@ class ClusterLoader(torch.utils.data.DataLoader):
 
     Args:
         cluster_data (torch_geometric.loader.ClusterData): The already
-            partioned data object.
+            partitioned data object.
         **kwargs (optional): Additional arguments of
             :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,
             :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.

torch_geometric/nn/conv/meshcnn_conv.py

@@ -64,7 +64,7 @@ class MeshCNNConv(MessagePassing):
         :math:`\mathcal{N}(1) = (a(1), b(1), c(1), d(1)) = (2, 3, 4, 5)`
 
 
-    Because of this ordering constrait, :obj:`MeshCNNConv` **requires
+    Because of this ordering constraint, :obj:`MeshCNNConv` **requires
     that the columns of** :math:`A`
     **be ordered in the following way**:
 
@@ -149,7 +149,7 @@ class MeshCNNConv(MessagePassing):
 
 
     Args:
-        in_channels (int): Corresonds to :math:`\text{Dim-Out}(k)`
+        in_channels (int): Corresponds to :math:`\text{Dim-Out}(k)`
             in the above overview. This
             represents the output dimension of the prior layer. For the given
             input mesh :math:`\mathcal{m} = (V, F)`, the prior layer is
@@ -184,7 +184,7 @@ class MeshCNNConv(MessagePassing):
             a vector of dimensions :attr:`out_channels`.
 
     Discussion:
-        The key difference that seperates :obj:`MeshCNNConv` from a traditional
+        The key difference that separates :obj:`MeshCNNConv` from a traditional
         message passing graph neural network is that :obj:`MeshCNNConv`
         requires the set of neighbors for a node
         :math:`\mathcal{N}(u) = (v_1, v_2, ...)`
@@ -198,7 +198,7 @@ class MeshCNNConv(MessagePassing):
         :math:`\mathbb{S}_4`. Put more plainly, in tradition message passing
         GNNs, the network is *unable* to distinguish one neighboring node
         from another.
-        In constrast, in :obj:`MeshCNNConv`, each of the 4 neighbors has a
+        In contrast, in :obj:`MeshCNNConv`, each of the 4 neighbors has a
         "role", either the "a", "b", "c", or "d" neighbor. We encode this fact
         by requiring that :math:`\mathcal{N}` return the 4-tuple,
         where the first component is the "a" neighbor, and so on.
@@ -444,7 +444,7 @@ class MeshCNNConv(MessagePassing):
         """
         assert isinstance(kernels, ModuleList), \
             f"Parameter 'kernels' must be a \
-            torch.nn.module.ModuleList with 5 memebers, but we got \
+            torch.nn.module.ModuleList with 5 members, but we got \
             {type(kernels)}."
 
         assert len(kernels) == 5, "Parameter 'kernels' must be a \

torch_geometric/nn/models/polynormer.py

@@ -37,7 +37,7 @@ class Polynormer(torch.nn.Module):
             (default: :obj:`True`)
         pre_ln (bool): Pre layer normalization.
             (default: :obj:`False`)
-        post_bn (bool): Post batch normlization.
+        post_bn (bool): Post batch normalization.
             (default: :obj:`True`)
         local_attn (bool): Whether use local attention.
             (default: :obj:`False`)

torch_geometric/nn/models/rev_gnn.py

@@ -196,8 +196,8 @@ class InvertibleModule(torch.nn.Module):
 class GroupAddRev(InvertibleModule):
     r"""The Grouped Reversible GNN module from the `"Graph Neural Networks with
     1000 Layers" <https://arxiv.org/abs/2106.07476>`_ paper.
-    This module enables training of arbitary deep GNNs with a memory complexity
-    independent of the number of layers.
+    This module enables training of arbitrary deep GNNs with a memory
+    complexity independent of the number of layers.
 
     It does so by partitioning input node features :math:`\mathbf{X}` into
     :math:`C` groups across the feature dimension. Then, a grouped reversible

torch_geometric/utils/influence.py

@@ -159,8 +159,8 @@ def jacobian_l1_agg_per_hop(
                             vectorize=vectorize)
     hop_subsets = k_hop_subsets_exact(node_idx, max_hops, edge_index,
                                       num_nodes, influence.device)
-    sigle_node_influence_per_hop = [influence[s].sum() for s in hop_subsets]
-    return torch.tensor(sigle_node_influence_per_hop, device=influence.device)
+    single_node_influence_per_hop = [influence[s].sum() for s in hop_subsets]
+    return torch.tensor(single_node_influence_per_hop, device=influence.device)
 
 
 def avg_total_influence(
@@ -169,7 +169,7 @@ def avg_total_influence(
 ) -> Tensor:
     """Compute the *influence‑weighted receptive field* ``R``."""
     avg_total_influences = torch.mean(influence_all_nodes, dim=0)
-    if normalize:  # nomalize by hop_0 (jacobian of the center node feature)
+    if normalize:  # normalize by hop_0 (jacobian of the center node feature)
         avg_total_influences = avg_total_influences / avg_total_influences[0]
     return avg_total_influences