pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/graphgym/imports.py

@@ -3,13 +3,24 @@ import warnings
 import torch
 
 try:
-    import pytorch_lightning as pl
+    import lightning.pytorch as pl
+    _pl_is_available = True
+except ImportError:
+    try:
+        import pytorch_lightning as pl
+        _pl_is_available = True
+    except ImportError:
+        _pl_is_available = False
+
+if _pl_is_available:
     LightningModule = pl.LightningModule
     Callback = pl.Callback
-except ImportError:
+else:
     pl = object
     LightningModule = torch.nn.Module
     Callback = object
 
-    warnings.warn("Please install 'pytorch_lightning' via  "
-                  "'pip install pytorch_lightning' in order to use GraphGym")
+    warnings.warn(
+        "To use GraphGym, install 'pytorch_lightning' or 'lightning' via "
+        "'pip install pytorch_lightning' or 'pip install lightning'",
+        stacklevel=2)

torch_geometric/graphgym/logger.py

@@ -239,7 +239,7 @@ def create_logger():
     r"""Create logger for the experiment."""
     loggers = []
     names = ['train', 'val', 'test']
-    for i, dataset in enumerate(range(cfg.share.num_splits)):
+    for i, _ in enumerate(range(cfg.share.num_splits)):
         loggers.append(Logger(name=names[i], task_type=infer_task()))
     return loggers
 

torch_geometric/graphgym/loss.py

@@ -10,7 +10,7 @@ def compute_loss(pred, true):
 
     Args:
         pred (torch.tensor): Unnormalized prediction
-        true (torch.tensor): Grou
+        true (torch.tensor): Ground truth labels
 
     Returns: Loss, normalized prediction score
 

torch_geometric/graphgym/models/encoder.py

@@ -53,7 +53,7 @@ class AtomEncoder(torch.nn.Module):
 
         self.atom_embedding_list = torch.nn.ModuleList()
 
-        for i, dim in enumerate(get_atom_feature_dims()):
+        for dim in get_atom_feature_dims():
             emb = torch.nn.Embedding(dim, emb_dim)
             torch.nn.init.xavier_uniform_(emb.weight.data)
             self.atom_embedding_list.append(emb)
@@ -87,7 +87,7 @@ class BondEncoder(torch.nn.Module):
 
         self.bond_embedding_list = torch.nn.ModuleList()
 
-        for i, dim in enumerate(get_bond_feature_dims()):
+        for dim in get_bond_feature_dims():
             emb = torch.nn.Embedding(dim, emb_dim)
             torch.nn.init.xavier_uniform_(emb.weight.data)
             self.bond_embedding_list.append(emb)

torch_geometric/graphgym/models/layer.py

@@ -52,7 +52,7 @@ def new_layer_config(
     has_bias: bool,
     cfg,
 ) -> LayerConfig:
-    r"""Createa a layer configuration for a GNN layer.
+    r"""Create a layer configuration for a GNN layer.
 
     Args:
         dim_in (int): The input feature dimension.

torch_geometric/graphgym/utils/comp_budget.py

@@ -18,8 +18,9 @@ def get_stats():
     return params_count(model)
 
 
-def match_computation(stats_baseline, key=['gnn', 'dim_inner'], mode='sqrt'):
+def match_computation(stats_baseline, key=None, mode='sqrt'):
     """Match computation budget by modifying :obj:`cfg.gnn.dim_inner`."""
+    key = key or ['gnn', 'dim_inner']
     stats = get_stats()
     if stats != stats_baseline:
         # Phase 1: fast approximation
@@ -67,12 +68,12 @@ def dict_to_stats(cfg_dict):
 
 def match_baseline_cfg(cfg_dict, cfg_dict_baseline, verbose=True):
     """Match the computational budget of a given baseline model. The current
-    configuration dictionary will be modifed and returned.
+    configuration dictionary will be modified and returned.
 
     Args:
         cfg_dict (dict): Current experiment's configuration
         cfg_dict_baseline (dict): Baseline configuration
-        verbose (str, optional): If printing matched paramter conunts
+        verbose (str, optional): If printing matched parameter conunts
     """
     from yacs.config import CfgNode as CN
     stats_baseline = dict_to_stats(cfg_dict_baseline)