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

torch_geometric/datasets/willow_object_class.py

@@ -142,7 +142,7 @@ class WILLOWObjectClass(InMemoryDataset):
             pos[:, 0] = pos[:, 0] * 256.0 / (img.size[0])
             pos[:, 1] = pos[:, 1] * 256.0 / (img.size[1])
 
-            img = img.resize((256, 256), resample=Image.BICUBIC)
+            img = img.resize((256, 256), resample=Image.Resampling.BICUBIC)
             img = transform(img)
 
             data = Data(img=img, pos=pos, name=name)

torch_geometric/datasets/word_net.py

@@ -67,7 +67,7 @@ class WordNet18(InMemoryDataset):
     def process(self) -> None:
         srcs, dsts, edge_types = [], [], []
         for path in self.raw_paths:
-            with open(path, 'r') as f:
+            with open(path) as f:
                 edges = [int(x) for x in f.read().split()[1:]]
                 edge = torch.tensor(edges, dtype=torch.long)
                 srcs.append(edge[::3])
@@ -173,7 +173,7 @@ class WordNet18RR(InMemoryDataset):
 
         srcs, dsts, edge_types = [], [], []
         for path in self.raw_paths:
-            with open(path, 'r') as f:
+            with open(path) as f:
                 edges = f.read().split()
 
                 _src = edges[::3]

torch_geometric/datasets/yelp.py

@@ -3,7 +3,6 @@ import os.path as osp
 from typing import Callable, List, Optional
 
 import numpy as np
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import Data, InMemoryDataset, download_google_url
@@ -73,6 +72,8 @@ class Yelp(InMemoryDataset):
         download_google_url(self.role_id, self.raw_dir, 'role.json')
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         f = np.load(osp.join(self.raw_dir, 'adj_full.npz'))
         adj = sp.csr_matrix((f['data'], f['indices'], f['indptr']), f['shape'])
         adj = adj.tocoo()

torch_geometric/datasets/zinc.py

@@ -139,7 +139,7 @@ class ZINC(InMemoryDataset):
             indices = list(range(len(mols)))
 
             if self.subset:
-                with open(osp.join(self.raw_dir, f'{split}.index'), 'r') as f:
+                with open(osp.join(self.raw_dir, f'{split}.index')) as f:
                     indices = [int(x) for x in f.read()[:-1].split(',')]
 
             pbar = tqdm(total=len(indices))

torch_geometric/device.py

@@ -0,0 +1,42 @@
+from typing import Any
+
+import torch
+
+
+def is_mps_available() -> bool:
+    r"""Returns a bool indicating if MPS is currently available."""
+    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        try:  # Github CI may not have access to MPS hardware. Confirm:
+            torch.empty(1, device='mps')
+            return True
+        except Exception:
+            return False
+    return False
+
+
+def is_xpu_available() -> bool:
+    r"""Returns a bool indicating if XPU is currently available."""
+    if hasattr(torch, 'xpu') and torch.xpu.is_available():
+        return True
+    try:
+        import intel_extension_for_pytorch as ipex
+        return ipex.xpu.is_available()
+    except ImportError:
+        return False
+
+
+def device(device: Any) -> torch.device:
+    r"""Returns a :class:`torch.device`.
+
+    If :obj:`"auto"` is specified, returns the optimal device depending on
+    available hardware.
+    """
+    if device != 'auto':
+        return torch.device(device)
+    if torch.cuda.is_available():
+        return torch.device('cuda')
+    if is_mps_available():
+        return torch.device('mps')
+    if is_xpu_available():
+        return torch.device('xpu')
+    return torch.device('cpu')

torch_geometric/distributed/local_feature_store.py

@@ -15,6 +15,7 @@ from torch_geometric.distributed.rpc import (
     rpc_async,
     rpc_register,
 )
+from torch_geometric.io import fs
 from torch_geometric.typing import EdgeType, NodeOrEdgeType, NodeType
 
 
@@ -415,11 +416,11 @@ class LocalFeatureStore(FeatureStore):
 
         node_feats: Optional[Dict[str, Any]] = None
         if osp.exists(osp.join(part_dir, 'node_feats.pt')):
-            node_feats = torch.load(osp.join(part_dir, 'node_feats.pt'))
+            node_feats = fs.torch_load(osp.join(part_dir, 'node_feats.pt'))
 
         edge_feats: Optional[Dict[str, Any]] = None
         if osp.exists(osp.join(part_dir, 'edge_feats.pt')):
-            edge_feats = torch.load(osp.join(part_dir, 'edge_feats.pt'))
+            edge_feats = fs.torch_load(osp.join(part_dir, 'edge_feats.pt'))
 
         if not meta['is_hetero'] and node_feats is not None:
             feat_store.put_global_id(node_feats['global_id'], group_name=None)

torch_geometric/distributed/local_graph_store.py

@@ -6,6 +6,7 @@ from torch import Tensor
 
 from torch_geometric.data import EdgeAttr, GraphStore
 from torch_geometric.distributed.partition import load_partition_info
+from torch_geometric.io import fs
 from torch_geometric.typing import EdgeTensorType, EdgeType, NodeType
 from torch_geometric.utils import sort_edge_index
 
@@ -185,7 +186,7 @@ class LocalGraphStore(GraphStore):
         graph_store.edge_pb = edge_pb
         graph_store.meta = meta
 
-        graph_data = torch.load(osp.join(part_dir, 'graph.pt'))
+        graph_data = fs.torch_load(osp.join(part_dir, 'graph.pt'))
         graph_store.is_sorted = meta['is_sorted']
 
         if not meta['is_hetero']:

torch_geometric/distributed/partition.py

@@ -3,15 +3,16 @@ import logging
 import os
 import os.path as osp
 from collections import defaultdict
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 
 import torch_geometric.distributed as pyg_dist
 from torch_geometric.data import Data, HeteroData
+from torch_geometric.io import fs
 from torch_geometric.loader.cluster import ClusterData
 from torch_geometric.sampler.utils import sort_csc
-from torch_geometric.typing import Dict, EdgeType, EdgeTypeStr, NodeType, Tuple
+from torch_geometric.typing import EdgeType, EdgeTypeStr, NodeType
 
 
 class Partitioner:
@@ -23,7 +24,7 @@ class Partitioner:
 
     **Homogeneous graphs:**
 
-    .. code-block::
+    .. code-block:: none
 
         root/
         |-- META.json
@@ -40,7 +41,7 @@ class Partitioner:
 
     **Heterogeneous graphs:**
 
-    .. code-block::
+    .. code-block:: none
 
         root/
         |-- META.json
@@ -380,21 +381,21 @@ def load_partition_info(
     assert osp.exists(partition_dir)
 
     if meta['is_hetero'] is False:
-        node_pb = torch.load(osp.join(root_dir, 'node_map.pt'))
-        edge_pb = torch.load(osp.join(root_dir, 'edge_map.pt'))
+        node_pb = fs.torch_load(osp.join(root_dir, 'node_map.pt'))
+        edge_pb = fs.torch_load(osp.join(root_dir, 'edge_map.pt'))
 
         return (meta, num_partitions, partition_idx, node_pb, edge_pb)
     else:
         node_pb_dict = {}
         node_pb_dir = osp.join(root_dir, 'node_map')
         for ntype in meta['node_types']:
-            node_pb_dict[ntype] = torch.load(
+            node_pb_dict[ntype] = fs.torch_load(
                 osp.join(node_pb_dir, f'{pyg_dist.utils.as_str(ntype)}.pt'))
 
         edge_pb_dict = {}
         edge_pb_dir = osp.join(root_dir, 'edge_map')
         for etype in meta['edge_types']:
-            edge_pb_dict[tuple(etype)] = torch.load(
+            edge_pb_dict[tuple(etype)] = fs.torch_load(
                 osp.join(edge_pb_dir, f'{pyg_dist.utils.as_str(etype)}.pt'))
 
         return (meta, num_partitions, partition_idx, node_pb_dict,