Add immutable transforms

CHANGELOG.md

@@ -45,6 +45,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- All transforms are now immutable, i.e., they perform a shallow-copy of the data and therefore do not longer modify data in-place ([#7429](https://github.com/pyg-team/pytorch_geometric/pull/7429))
 - Set `output_size` in the `repeat_interleave` operation in `QuantileAggregation` ([#7426](https://github.com/pyg-team/pytorch_geometric/pull/7426))
 - Fixed gradient computation of edge weights in `utils.spmm` ([#7428](https://github.com/pyg-team/pytorch_geometric/pull/7428))
 - Re-factored `ClusterLoader` to integrate `pyg-lib` METIS routine ([#7416](https://github.com/pyg-team/pytorch_geometric/pull/7416))

examples/colors_topk_pool.py

@@ -1,3 +1,4 @@
+import copy
 import os.path as osp
 
 import torch
@@ -14,6 +15,7 @@
 
 class HandleNodeAttention:
     def __call__(self, data):
+        data = copy.copy(data)
         data.attn = torch.softmax(data.x[:, 0], dim=0)
         data.x = data.x[:, 1:]
         return data

examples/proteins_diff_pool.py

@@ -11,16 +11,14 @@
 
 max_nodes = 150
 
-
-class MyFilter:
-    def __call__(self, data):
-        return data.num_nodes <= max_nodes
-
-
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data',
                 'PROTEINS_dense')
-dataset = TUDataset(path, name='PROTEINS', transform=T.ToDense(max_nodes),
-                    pre_filter=MyFilter())
+dataset = TUDataset(
+    path,
+    name='PROTEINS',
+    transform=T.ToDense(max_nodes),
+    pre_filter=lambda data: data.num_nodes <= max_nodes,
+)
 dataset = dataset.shuffle()
 n = (len(dataset) + 9) // 10
 test_dataset = dataset[:n]

examples/qm9_nn_conv.py

@@ -1,3 +1,4 @@
+import copy
 import os.path as osp
 
 import torch
@@ -16,13 +17,14 @@
 
 class MyTransform:
     def __call__(self, data):
-        # Specify target.
-        data.y = data.y[:, target]
+        data = copy.copy(data)
+        data.y = data.y[:, target]  # Specify target.
         return data
 
 
 class Complete:
     def __call__(self, data):
+        data = copy.copy(data)
         device = data.edge_index.device
 
         row = torch.arange(data.num_nodes, dtype=torch.long, device=device)

examples/triangles_sag_pool.py

@@ -1,3 +1,4 @@
+import copy
 import os.path as osp
 
 import torch
@@ -15,6 +16,7 @@
 
 class HandleNodeAttention:
     def __call__(self, data):
+        data = copy.copy(data)
         data.attn = torch.softmax(data.x, dim=0).flatten()
         data.x = None
         return data

torch_geometric/data/in_memory_dataset.py

@@ -85,6 +85,7 @@ def len(self) -> int:
         return 0
 
     def get(self, idx: int) -> Data:
+        # TODO (matthias) Avoid unnecessary copy here.
         if self.len() == 1:
             return copy.copy(self._data)
 

torch_geometric/transforms/add_metapaths.py

@@ -129,7 +129,7 @@ def __init__(
         self.max_sample = max_sample
         self.weighted = weighted
 
-    def __call__(self, data: HeteroData) -> HeteroData:
+    def forward(self, data: HeteroData) -> HeteroData:
         edge_types = data.edge_types  # save original edge types
         data.metapath_dict = {}
 
@@ -244,7 +244,7 @@ def __init__(
         assert len(walks_per_node) == len(metapaths)
         self.walks_per_node = walks_per_node
 
-    def __call__(self, data: HeteroData) -> HeteroData:
+    def forward(self, data: HeteroData) -> HeteroData:
         edge_types = data.edge_types  # save original edge types
         data.metapath_dict = {}
 

torch_geometric/transforms/add_positional_encoding.py

@@ -64,7 +64,7 @@ def __init__(
         self.is_undirected = is_undirected
         self.kwargs = kwargs
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         from scipy.sparse.linalg import eigs, eigsh
         eig_fn = eigs if not self.is_undirected else eigsh
 
@@ -117,7 +117,7 @@ def __init__(
         self.walk_length = walk_length
         self.attr_name = attr_name
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         row, col = data.edge_index
         N = data.num_nodes
 

torch_geometric/transforms/add_remaining_self_loops.py

@@ -32,7 +32,7 @@ def __init__(self, attr: Optional[str] = 'edge_weight',
         self.attr = attr
         self.fill_value = fill_value
 
-    def __call__(
+    def forward(
         self,
         data: Union[Data, HeteroData],
     ) -> Union[Data, HeteroData]:

torch_geometric/transforms/add_self_loops.py

@@ -32,7 +32,7 @@ def __init__(self, attr: Optional[str] = 'edge_weight',
         self.attr = attr
         self.fill_value = fill_value
 
-    def __call__(
+    def forward(
         self,
         data: Union[Data, HeteroData],
     ) -> Union[Data, HeteroData]:

torch_geometric/transforms/base_transform.py

@@ -1,4 +1,5 @@
-from abc import ABC
+import copy
+from abc import ABC, abstractmethod
 from typing import Any
 
 
@@ -27,7 +28,11 @@ class BaseTransform(ABC):
         data = transform(data)  # Explicitly transform data.
     """
     def __call__(self, data: Any) -> Any:
-        raise NotImplementedError
+        return self.forward(copy.copy(data))
+
+    @abstractmethod
+    def forward(self, data: Any) -> Any:
+        pass
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}()'

torch_geometric/transforms/cartesian.py

@@ -32,7 +32,7 @@ def __init__(
         self.max = max_value
         self.cat = cat
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
 
         cart = pos[row] - pos[col]

torch_geometric/transforms/center.py

@@ -9,7 +9,7 @@
 class Center(BaseTransform):
     r"""Centers node positions :obj:`data.pos` around the origin
     (functional name: :obj:`center`)."""
-    def __call__(
+    def forward(
         self,
         data: Union[Data, HeteroData],
     ) -> Union[Data, HeteroData]:

torch_geometric/transforms/compose.py

@@ -13,7 +13,7 @@ class Compose(BaseTransform):
     def __init__(self, transforms: List[Callable]):
         self.transforms = transforms
 
-    def __call__(
+    def forward(
         self,
         data: Union[Data, HeteroData],
     ) -> Union[Data, HeteroData]:

torch_geometric/transforms/constant.py

@@ -34,7 +34,7 @@ def __init__(
         self.cat = cat
         self.node_types = node_types
 
-    def __call__(
+    def forward(
         self,
         data: Union[Data, HeteroData],
     ) -> Union[Data, HeteroData]:

torch_geometric/transforms/delaunay.py

@@ -10,7 +10,7 @@
 class Delaunay(BaseTransform):
     r"""Computes the delaunay triangulation of a set of points
     (functional name: :obj:`delaunay`)."""
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         if data.pos.size(0) < 2:
             data.edge_index = torch.tensor([], dtype=torch.long,
                                            device=data.pos.device).view(2, 0)

torch_geometric/transforms/distance.py

@@ -27,7 +27,7 @@ def __init__(self, norm: bool = True, max_value: Optional[float] = None,
         self.max = max_value
         self.cat = cat
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
 
         dist = torch.norm(pos[col] - pos[row], p=2, dim=-1).view(-1, 1)

torch_geometric/transforms/face_to_edge.py

@@ -18,7 +18,7 @@ class FaceToEdge(BaseTransform):
     def __init__(self, remove_faces: bool = True):
         self.remove_faces = remove_faces
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         if hasattr(data, 'face'):
             face = data.face
             edge_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)

torch_geometric/transforms/feature_propagation.py

@@ -40,7 +40,7 @@ def __init__(self, missing_mask: Tensor, num_iterations: int = 40):
         self.missing_mask = missing_mask
         self.num_iterations = num_iterations
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         assert 'edge_index' in data or 'adj_t' in data
         assert data.x.size() == self.missing_mask.size()
 

torch_geometric/transforms/fixed_points.py

@@ -37,7 +37,7 @@ def __init__(
         self.replace = replace
         self.allow_duplicates = allow_duplicates
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         num_nodes = data.num_nodes
 
         if self.replace:

torch_geometric/transforms/gcn_norm.py

@@ -19,7 +19,7 @@ class GCNNorm(BaseTransform):
     def __init__(self, add_self_loops: bool = True):
         self.add_self_loops = add_self_loops
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         gcn_norm = torch_geometric.nn.conv.gcn_conv.gcn_norm
         assert 'edge_index' in data or 'adj_t' in data
 

torch_geometric/transforms/gdc.py

@@ -97,7 +97,7 @@ def __init__(
             assert exact or self_loop_weight == 1
 
     @torch.no_grad()
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         N = data.num_nodes
         edge_index = data.edge_index
         if data.edge_attr is None:

torch_geometric/transforms/generate_mesh_normals.py

@@ -11,7 +11,7 @@
 class GenerateMeshNormals(BaseTransform):
     r"""Generate normal vectors for each mesh node based on neighboring
     faces (functional name: :obj:`generate_mesh_normals`)."""
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         assert 'face' in data
         pos, face = data.pos, data.face
 

torch_geometric/transforms/grid_sampling.py

@@ -36,7 +36,7 @@ def __init__(self, size: Union[float, List[float], Tensor],
         self.start = start
         self.end = end
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         num_nodes = data.num_nodes
 
         batch = data.get('batch', None)

torch_geometric/transforms/knn_graph.py

@@ -45,7 +45,7 @@ def __init__(
         self.cosine = cosine
         self.num_workers = num_workers
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         data.edge_attr = None
         batch = data.batch if 'batch' in data else None
 

torch_geometric/transforms/laplacian_lambda_max.py

@@ -40,7 +40,7 @@ def __init__(
         self.normalization = normalization
         self.is_undirected = is_undirected
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         edge_weight = data.edge_attr
         if edge_weight is not None and edge_weight.numel() != data.num_edges:
             edge_weight = None

torch_geometric/transforms/largest_connected_components.py

@@ -28,7 +28,7 @@ def __init__(self, num_components: int = 1, connection: str = 'weak'):
         self.num_components = num_components
         self.connection = connection
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         import numpy as np
         import scipy.sparse as sp
 

torch_geometric/transforms/line_graph.py

@@ -34,7 +34,7 @@ class LineGraph(BaseTransform):
     def __init__(self, force_directed: bool = False):
         self.force_directed = force_directed
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         N = data.num_nodes
         edge_index, edge_attr = data.edge_index, data.edge_attr
         edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes=N)

torch_geometric/transforms/linear_transformation.py

@@ -26,10 +26,10 @@ def __init__(self, matrix: Tensor):
             f'Transformation matrix should be square (got {matrix.size()})')
 
         # Store the matrix as its transpose.
-        # We do this to enable post-multiplication in `__call__`.
+        # We do this to enable post-multiplication in `forward`.
         self.matrix = matrix.t()
 
-    def __call__(
+    def forward(
         self,
         data: Union[Data, HeteroData],
     ) -> Union[Data, HeteroData]:

torch_geometric/transforms/local_cartesian.py

@@ -23,7 +23,7 @@ def __init__(self, norm: bool = True, cat: bool = True):
         self.norm = norm
         self.cat = cat
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr
 
         cart = pos[row] - pos[col]

torch_geometric/transforms/local_degree_profile.py

@@ -24,7 +24,7 @@ def __init__(self):
         from torch_geometric.nn.aggr.fused import FusedAggregation
         self.aggr = FusedAggregation(['min', 'max', 'mean', 'std'])
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         row, col = data.edge_index
         N = data.num_nodes
 

torch_geometric/transforms/mask.py

@@ -52,7 +52,7 @@ def __init__(
         self.sizes = sizes
         self.replace = replace
 
-    def __call__(
+    def forward(
         self,
         data: Union[Data, HeteroData],
     ) -> Union[Data, HeteroData]:
@@ -106,7 +106,7 @@ def __init__(
         self.attrs = [attrs] if isinstance(attrs, str) else attrs
         self.replace = replace
 
-    def __call__(
+    def forward(
         self,
         data: Union[Data, HeteroData],
     ) -> Union[Data, HeteroData]:

torch_geometric/transforms/node_property_split.py

@@ -81,8 +81,7 @@ def __init__(
         self.ratios = ratios
         self.ascending = ascending
 
-    def __call__(self, data: Data) -> Data:
-
+    def forward(self, data: Data) -> Data:
         G = to_networkx(data, to_undirected=True, remove_self_loops=True)
         property_values = self.compute_fn(G, self.ascending)
         mask_dict = self._mask_nodes_by_property(property_values, self.ratios)

torch_geometric/transforms/normalize_features.py

@@ -17,7 +17,7 @@ class NormalizeFeatures(BaseTransform):
     def __init__(self, attrs: List[str] = ["x"]):
         self.attrs = attrs
 
-    def __call__(
+    def forward(
         self,
         data: Union[Data, HeteroData],
     ) -> Union[Data, HeteroData]:

torch_geometric/transforms/normalize_rotation.py

@@ -24,7 +24,7 @@ def __init__(self, max_points: int = -1, sort: bool = False):
         self.max_points = max_points
         self.sort = sort
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         pos = data.pos
 
         if self.max_points > 0 and pos.size(0) > self.max_points:

torch_geometric/transforms/normalize_scale.py

@@ -11,7 +11,7 @@ class NormalizeScale(BaseTransform):
     def __init__(self):
         self.center = Center()
 
-    def __call__(self, data: Data) -> Data:
+    def forward(self, data: Data) -> Data:
         data = self.center(data)
 
         scale = (1 / data.pos.abs().max()) * 0.999999