pyg-nightly 2.6.0.dev20240511__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/data/hetero_data.py

@@ -337,7 +337,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
     def __cat_dim__(self, key: str, value: Any,
                     store: Optional[NodeOrEdgeStorage] = None, *args,
                     **kwargs) -> Any:
-        if is_sparse(value) and 'adj' in key:
+        if is_sparse(value) and ('adj' in key or 'edge_index' in key):
             return (0, 1)
         elif isinstance(store, EdgeStorage) and 'index' in key:
             return -1
@@ -780,8 +780,8 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         for edge_type in self.edge_types:
             if edge_type not in edge_types:
                 del data[edge_type]
-        node_types = set(e[0] for e in edge_types)
-        node_types |= set(e[-1] for e in edge_types)
+        node_types = {e[0] for e in edge_types}
+        node_types |= {e[-1] for e in edge_types}
         for node_type in self.node_types:
             if node_type not in node_types:
                 del data[node_type]
@@ -887,7 +887,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
                 if len(sizes) != len(stores):
                     continue
                 # The attributes needs to have the same number of dimensions:
-                lengths = set([len(size) for size in sizes])
+                lengths = {len(size) for size in sizes}
                 if len(lengths) != 1:
                     continue
                 # The attributes needs to have the same size in all dimensions:

torch_geometric/data/in_memory_dataset.py

@@ -347,10 +347,8 @@ class InMemoryDataset(Dataset):
 def nested_iter(node: Union[Mapping, Sequence]) -> Iterable:
     if isinstance(node, Mapping):
         for key, value in node.items():
-            for inner_key, inner_value in nested_iter(value):
-                yield inner_key, inner_value
+            yield from nested_iter(value)
     elif isinstance(node, Sequence):
-        for i, inner_value in enumerate(node):
-            yield i, inner_value
+        yield from enumerate(node)
     else:
         yield None, node