pyg-nightly 2.7.0.dev20250822__py3-none-any.whl → 2.7.0.dev20250824__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250822.dist-info → pyg_nightly-2.7.0.dev20250824.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=EzPtCTlUiMKhCY8DpZ-PoeX7k2ZTyBUGvpb8T5aKXFg,2250
+torch_geometric/__init__.py,sha256=MgxbhqEQw85B1XawAzP1xK2jTO9z_QcIYtwhfU88AQQ,2250
 torch_geometric/_compile.py,sha256=9yqMTBKatZPr40WavJz9FjNi7pQj8YZAZOyZmmRGXgc,1351
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -519,10 +519,10 @@ torch_geometric/profile/profile.py,sha256=cHCY4U0XtyqyKC5u380q6TspsOZ5tGHNXaZsKu
 torch_geometric/profile/profiler.py,sha256=rfNciRzWDka_BgO6aPFi3cy8mcT4lSgFWy-WfPgI2SI,16891
 torch_geometric/profile/utils.py,sha256=ynlUVemNJZ6XjJKIkPJNwFPoUyBgVAnchfHBpKOp_HE,5903
 torch_geometric/sampler/__init__.py,sha256=0h_xJ7CQnlTxF5hUpc81WPQ0QaBtouG8eKK1RzPGA-s,512
-torch_geometric/sampler/base.py,sha256=T7RMx14RSlEKlQUkMvR1EzREaXi14VgR5GIwLfvbXzQ,27055
+torch_geometric/sampler/base.py,sha256=0z5e6DYqUGa1-eyJKUGCyUJe6T-DoKBVpOdNxD-dOzM,29225
 torch_geometric/sampler/hgt_sampler.py,sha256=jizRJyEoz4WBOEELuqdytG2hB3UpVQX7yVPM83kvpfE,2991
 torch_geometric/sampler/neighbor_sampler.py,sha256=GBnIwPF5pf85wDo_FaMKtMgBlQhzRx5ysiAomMNhCrQ,34092
-torch_geometric/sampler/utils.py,sha256=RJtasO6Q7Pp3oYEOWrbf2DEYuSfuKZOsF2I7-eJDnoA,5485
+torch_geometric/sampler/utils.py,sha256=U4XGCEBX6rhmKHpCwYjskNvKNnRTXb9GeVKR1vSWuvA,7250
 torch_geometric/testing/__init__.py,sha256=m3yp_5UnCAxVgzTFofpiVt0vdbl5GwVAve8WTrAaNxo,1319
 torch_geometric/testing/asserts.py,sha256=DLC9HnBgFWuTIiQs2OalsQcXGhOVG-e6R99IWhkO32c,4606
 torch_geometric/testing/data.py,sha256=O1qo8FyNxt6RGf63Ys3eXBfa5RvYydeZLk74szrez3c,2604
@@ -645,7 +645,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.dev20250822.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250822.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250822.dist-info/METADATA,sha256=dM0G4qt1GnfQIJxN5QRCrRemxyP3r1iyITMakFRVMAw,64100
-pyg_nightly-2.7.0.dev20250822.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250824.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250824.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250824.dist-info/METADATA,sha256=e414Fyh9TNPZurzgwbAnQfQufV29PLugLmumo4GOIyk,64100
+pyg_nightly-2.7.0.dev20250824.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.dev20250822'
+__version__ = '2.7.0.dev20250824'
 
 __all__ = [
     'Index',

torch_geometric/sampler/base.py

@@ -11,7 +11,10 @@ import torch
 from torch import Tensor
 
 from torch_geometric.data import Data, FeatureStore, GraphStore, HeteroData
-from torch_geometric.sampler.utils import to_bidirectional
+from torch_geometric.sampler.utils import (
+    local_to_global_node_idx,
+    to_bidirectional,
+)
 from torch_geometric.typing import EdgeType, EdgeTypeStr, NodeType, OptTensor
 from torch_geometric.utils.mixin import CastMixin
 
@@ -207,6 +210,29 @@ class SamplerOutput(CastMixin):
     # API for the expected output of a sampler.
     metadata: Optional[Any] = None
 
+    @property
+    def global_row(self) -> Tensor:
+        return local_to_global_node_idx(self.node, self.row)
+
+    @property
+    def global_col(self) -> Tensor:
+        return local_to_global_node_idx(self.node, self.col)
+
+    @property
+    def seed_node(self) -> Tensor:
+        return local_to_global_node_idx(
+            self.node, self.batch) if self.batch is not None else None
+
+    @property
+    def global_orig_row(self) -> Tensor:
+        return local_to_global_node_idx(
+            self.node, self.orig_row) if self.orig_row is not None else None
+
+    @property
+    def global_orig_col(self) -> Tensor:
+        return local_to_global_node_idx(
+            self.node, self.orig_col) if self.orig_col is not None else None
+
     def to_bidirectional(
         self,
         keep_orig_edges: bool = False,
@@ -294,6 +320,43 @@ class HeteroSamplerOutput(CastMixin):
     # API for the expected output of a sampler.
     metadata: Optional[Any] = None
 
+    @property
+    def global_row(self) -> Dict[EdgeType, Tensor]:
+        return {
+            edge_type: local_to_global_node_idx(self.node[edge_type[0]], row)
+            for edge_type, row in self.row.items()
+        }
+
+    @property
+    def global_col(self) -> Dict[EdgeType, Tensor]:
+        return {
+            edge_type: local_to_global_node_idx(self.node[edge_type[2]], col)
+            for edge_type, col in self.col.items()
+        }
+
+    @property
+    def seed_node(self) -> Optional[Dict[NodeType, Tensor]]:
+        return {
+            node_type: local_to_global_node_idx(self.node[node_type], batch)
+            for node_type, batch in self.batch.items()
+        } if self.batch is not None else None
+
+    @property
+    def global_orig_row(self) -> Optional[Dict[EdgeType, Tensor]]:
+        return {
+            edge_type: local_to_global_node_idx(self.node[edge_type[0]],
+                                                orig_row)
+            for edge_type, orig_row in self.orig_row.items()
+        } if self.orig_row is not None else None
+
+    @property
+    def global_orig_col(self) -> Optional[Dict[EdgeType, Tensor]]:
+        return {
+            edge_type: local_to_global_node_idx(self.node[edge_type[2]],
+                                                orig_col)
+            for edge_type, orig_col in self.orig_col.items()
+        } if self.orig_col is not None else None
+
     def to_bidirectional(
         self,
         keep_orig_edges: bool = False,

torch_geometric/sampler/utils.py

@@ -160,3 +160,48 @@ def remap_keys(
         k if k in exclude else mapping.get(k, k): v
         for k, v in inputs.items()
     }
+
+
+def local_to_global_node_idx(node_values: Tensor,
+                             local_indices: Tensor) -> Tensor:
+    """Convert a tensor of indices referring to elements in the node_values
+    tensor to their values.
+
+    Args:
+        node_values (Tensor): The node values. (num_nodes, feature_dim)
+        local_indices (Tensor): The local indices. (num_indices)
+
+    Returns:
+        Tensor: The values of the node_values tensor at the local indices.
+        (num_indices, feature_dim)
+    """
+    return torch.index_select(node_values, dim=0, index=local_indices)
+
+
+def global_to_local_node_idx(node_values: Tensor,
+                             local_values: Tensor) -> Tensor:
+    """Converts a tensor of values that are contained in the node_values
+    tensor to their indices in that tensor.
+
+    Args:
+        node_values (Tensor): The node values. (num_nodes, feature_dim)
+        local_values (Tensor): The local values. (num_indices, feature_dim)
+
+    Returns:
+        Tensor: The indices of the local values in the node_values tensor.
+        (num_indices)
+    """
+    if node_values.dim() == 1:
+        node_values = node_values.unsqueeze(1)
+    if local_values.dim() == 1:
+        local_values = local_values.unsqueeze(1)
+    node_values_expand = node_values.unsqueeze(-1).expand(
+        *node_values.shape,
+        local_values.shape[0])  # (num_nodes, feature_dim, num_indices)
+    local_values_expand = local_values.transpose(0, 1).unsqueeze(0).expand(
+        *node_values_expand.shape)  # (num_nodes, feature_dim, num_indices)
+    idx_match = torch.all(node_values_expand == local_values_expand,
+                          dim=1).nonzero()  # (num_indices, 2)
+    sort_idx = torch.argsort(idx_match[:, 1])
+
+    return idx_match[:, 0][sort_idx]