Repository: pyg-team/pytorch_geometric (Python)

Issue:
CI Failure: `test_lr_scheduler_resolver[scheduler_args5]` & `test_add_weighted_metapaths`
### 🐛 Describe the bug

```
====================================================================================================== FAILURES ======================================================================================================
____________________________________________________________________________________ test_lr_scheduler_resolver[scheduler_args5] _____________________________________________________________________________________

scheduler_args = ('constant', <class 'torch.optim.lr_scheduler.ConstantLR'>)

    @pytest.mark.parametrize('scheduler_args', [
        ("constant_with_warmup", LambdaLR),
        ("linear_with_warmup", LambdaLR),
        ("cosine_with_warmup", LambdaLR),
        ("cosine_with_warmup_restarts", LambdaLR),
        ("polynomial_with_warmup", LambdaLR),
        ("constant", ConstantLR),
        ('ReduceLROnPlateau', ReduceLROnPlateau),
    ])
    def test_lr_scheduler_resolver(scheduler_args):
        scheduler_name, scheduler_cls = scheduler_args
    
        model = torch.nn.Linear(10, 5)
        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
>       lr_scheduler = lr_scheduler_resolver(
            scheduler_name,
            optimizer,
            num_training_steps=100,
        )

test/nn/test_resolver.py:76: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

query = 'constant'
optimizer = Adam (
Parameter Group 0
    amsgrad: False
    betas: (0.9, 0.999)
    capturable: False
    differentiable: False
    eps: 1e-08
    foreach: None
    fused: False
    lr: 0.01
    maximize: False
    weight_decay: 0
)
warmup_ratio_or_steps = 0.1, num_training_steps = 100, kwargs = {}, warmup_steps = 10
classes = [<class 'torch.optim.lr_scheduler._LRScheduler'>, <class 'torch.optim.lr_scheduler.ReduceLROnPlateau'>, <class 'torch_...metric.nn.lr_scheduler.CosineWithWarmupLR'>, <class 'torch_geometric.nn.lr_scheduler.CosineWithWarmupRestartsLR'>, ...]
customized_lr_schedulers = [<class 'torch_geometric.nn.lr_scheduler.ConstantWithWarmupLR'>, <class 'torch_geometric.nn.lr_scheduler.LinearWithWar...ometric.nn.lr_scheduler.CosineWithWarmupRestartsLR'>, <class 'torch_geometric.nn.lr_scheduler.PolynomialWithWarmupLR'>]
query_repr = 'constant', base_cls_repr = 'lr', cls = <class 'torch_geometric.nn.lr_scheduler.PolynomialWithWarmupLR'>, cls_repr = 'polynomialwithwarmuplr'

    def lr_scheduler_resolver(
        query: Union[Any, str],
        optimizer: Optimizer,
        warmup_ratio_or_steps: Optional[Union[float, int]] = 0.1,
        num_training_steps: Optional[int] = None,
        **kwargs,
    ) -> Union[_LRScheduler, ReduceLROnPlateau]:
        r"""A resolver to obtain a learning rate scheduler implemented in either
        PyG or PyTorch from its name or type.
    
        Args:
            query (Any or str): The query name of the learning rate scheduler.
            optimizer (Optimizer): The optimizer to be scheduled.
            warmup_ratio_or_steps (float or int, optional): The number of warmup
                steps. If given as a `float`, it will act as a ratio that gets
                multiplied with the number of training steps to obtain the number
                of warmup steps. Only required for warmup-based LR schedulers.
                (default: :obj:`0.1`)
            num_training_steps (int, optional): The total number of training steps.
                (default: :obj:`None`)
            **kwargs (optional): Additional arguments of the LR scheduler.
        """
        if not isinstance(query, str):
            return query
    
        if isinstance(warmup_ratio_or_steps, float):
            if warmup_ratio_or_steps < 0 or warmup_ratio_or_steps > 1:
                raise ValueError(f"`warmup_ratio_or_steps` needs to be between "
                                 f"0.0 and 1.0 when given as a floating point "
                                 f"number (got {warmup_ratio_or_steps}).")
            if num_training_steps is not None:
                warmup_steps = round(warmup_ratio_or_steps * num_training_steps)
        elif isinstance(warmup_ratio_or_steps, int):
            if warmup_ratio_or_steps < 0:
                raise ValueError(f"`warmup_ratio_or_steps` needs to be positive "
                                 f"when given as an integer "
                                 f"(got {warmup_ratio_or_steps}).")
            warmup_steps = warmup_ratio_or_steps
        else:
            raise ValueError(f"Found invalid type of `warmup_ratio_or_steps` "
                             f"(got {type(warmup_ratio_or_steps)})")
    
        base_cls = _LRScheduler
        classes = [
            scheduler for scheduler in vars(torch.optim.lr_scheduler).values()
            if isinstance(scheduler, type) and issubclass(scheduler, base_cls)
        ] + [ReduceLROnPlateau]
    
        customized_lr_schedulers = [
            ConstantWithWarmupLR,
            LinearWithWarmupLR,
            CosineWithWarmupLR,
            CosineWithWarmupRestartsLR,
            PolynomialWithWarmupLR,
        ]
        classes += customized_lr_schedulers
    
        query_repr = normalize_string(query)
        base_cls_repr = normalize_string('LR')
    
        for cls in classes:
            cls_repr = normalize_string(cls.__name__)
            if query_repr in [cls_repr, cls_repr.replace(base_cls_repr, '')]:
                if inspect.isclass(cls):
                    if cls in customized_lr_schedulers:
                        cls_keys = inspect.signature(cls).parameters.keys()
                        if 'num_warmup_steps' in cls_keys:
                            kwargs['num_warmup_steps'] = warmup_steps
                        if 'num_training_steps' in cls_keys:
                            kwargs['num_training_steps'] = num_training_steps
                    obj = cls(optimizer, **kwargs)
                    return obj
                return cls
    
        choices = set(cls.__name__ for cls in classes)
>       raise ValueError(f"Could not resolve '{query}' among choices {choices}")
E       ValueError: Could not resolve 'constant' among choices {'ReduceLROnPlateau', 'PolynomialWithWarmupLR', '_LRScheduler', 'LinearWithWarmupLR', 'CosineWithWarmupRestartsLR', 'CosineWithWarmupLR', 'ConstantWithWarmupLR'}

/usr/local/lib/python3.8/dist-packages/torch_geometric-2.2.0-py3.8.egg/torch_geometric/nn/resolver.py:192: ValueError
____________________________________________________________________________________________ test_add_weighted_metapaths _____________________________________________________________________________________________

    def test_add_weighted_metapaths():
        torch.manual_seed(12345)
    
        data = HeteroData()
        data['a'].num_nodes = 2
        data['b'].num_nodes = 3
        data['c'].num_nodes = 2
        data['d'].num_nodes = 2
        data['a', 'b'].edge_index = tensor([[0, 1, 1], [0, 1, 2]])
        data['b', 'a'].edge_index = data['a', 'b'].edge_index.flip([0])
        data['b', 'c'].edge_index = tensor([[0, 1, 2], [0, 1, 1]])
        data['c', 'b'].edge_index = data['b', 'c'].edge_index.flip([0])
        data['c', 'd'].edge_index = tensor([[0, 1], [0, 0]])
        data['d', 'c'].edge_index = data['c', 'd'].edge_index.flip([0])
    
        metapaths = [
            [('a', 'b'), ('b', 'c')],
            [('a', 'b'), ('b', 'c'), ('c', 'd')],
            [('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'c'), ('c', 'b'),
             ('b', 'a')],
        ]
        transform = AddMetaPaths(metapaths, weighted=True)
        metapath_data = transform(copy.copy(data))
    
        # Make sure manually added metapaths compute the correct number of edges
        assert metapath_data['a', 'c'].edge_weight.tolist() == [1, 2]
        assert metapath_data['a', 'd'].edge_weight.tolist() == [1, 2]
>       assert metapath_data['a', 'a'].edge_weight.tolist() == [1, 2, 2, 4]
E       assert [2.0, 1.0, 4.0, 2.0] == [1, 2, 2, 4]
E         At index 0 diff: 2.0 != 1
E         Use -v to get more diff

test/transforms/test_add_metapaths.py:115: AssertionError
```

### Environment

* Versions: 
torch                   1.14.0a0+44dac51
torch-cluster           1.6.0
torch-geometric         2.2.0
torch-scatter           2.1.0
torch-sparse            0.6.16
torch-spline-conv       1.2.1
* OS: Ubuntu
* Python version: 3.8.10
* CUDA 12
* How you installed PyTorch and PyG (`conda`, `pip`, source): source



Discussion:
Just to understand these issues: Is this ran on CPU or on CUDA?
This should be fixed with https://github.com/pyg-team/pytorch_geometric/pull/6353. Please let me know if it is working.



Patch:
diff --git a/torch_geometric/nn/resolver.py b/torch_geometric/nn/resolver.py
index 133a971e3eab..ce6a9495848d 100644
--- a/torch_geometric/nn/resolver.py
+++ b/torch_geometric/nn/resolver.py
@@ -4,7 +4,7 @@
 import torch
 from torch import Tensor
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import ReduceLROnPlateau, _LRScheduler
+from torch.optim.lr_scheduler import ReduceLROnPlateau
 
 from torch_geometric.nn.lr_scheduler import (
     ConstantWithWarmupLR,
@@ -15,6 +15,11 @@
 )
 from torch_geometric.resolver import normalize_string, resolver
 
+try:
+    from torch.optim.lr_scheduler import LRScheduler
+except ImportError:  # PyTorch < 2.0
+    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+
 # Activation Resolver #########################################################
 
 
@@ -81,7 +86,7 @@ def lr_scheduler_resolver(
     warmup_ratio_or_steps: Optional[Union[float, int]] = 0.1,
     num_training_steps: Optional[int] = None,
     **kwargs,
-) -> Union[_LRScheduler, ReduceLROnPlateau]:
+) -> Union[LRScheduler, ReduceLROnPlateau]:
     r"""A resolver to obtain a learning rate scheduler implemented in either
     PyG or PyTorch from its name or type.
 
@@ -117,7 +122,7 @@ def lr_scheduler_resolver(
         raise ValueError(f"Found invalid type of `warmup_ratio_or_steps` "
                          f"(got {type(warmup_ratio_or_steps)})")
 
-    base_cls = _LRScheduler
+    base_cls = LRScheduler
     classes = [
         scheduler for scheduler in vars(torch.optim.lr_scheduler).values()
         if isinstance(scheduler, type) and issubclass(scheduler, base_cls)


Test patch:
diff --git a/test/nn/test_resolver.py b/test/nn/test_resolver.py
index a1e57f244664..46ffd4ed8dcd 100644
--- a/test/nn/test_resolver.py
+++ b/test/nn/test_resolver.py
@@ -59,12 +59,12 @@ def test_normalization_resolver(norm_tuple):
 
 
 @pytest.mark.parametrize('scheduler_args', [
-    ("constant_with_warmup", LambdaLR),
-    ("linear_with_warmup", LambdaLR),
-    ("cosine_with_warmup", LambdaLR),
-    ("cosine_with_warmup_restarts", LambdaLR),
-    ("polynomial_with_warmup", LambdaLR),
-    ("constant", ConstantLR),
+    ('constant_with_warmup', LambdaLR),
+    ('linear_with_warmup', LambdaLR),
+    ('cosine_with_warmup', LambdaLR),
+    ('cosine_with_warmup_restarts', LambdaLR),
+    ('polynomial_with_warmup', LambdaLR),
+    ('constant', ConstantLR),
     ('ReduceLROnPlateau', ReduceLROnPlateau),
 ])
 def test_lr_scheduler_resolver(scheduler_args):
diff --git a/test/transforms/test_add_metapaths.py b/test/transforms/test_add_metapaths.py
index aa97df044c16..a3f2403bbc44 100644
--- a/test/transforms/test_add_metapaths.py
+++ b/test/transforms/test_add_metapaths.py
@@ -5,6 +5,7 @@
 
 from torch_geometric.data import HeteroData
 from torch_geometric.transforms import AddMetaPaths, AddRandomMetaPaths
+from torch_geometric.utils import coalesce
 
 
 def generate_data() -> HeteroData:
@@ -107,25 +108,40 @@ def test_add_weighted_metapaths():
          ('b', 'a')],
     ]
     transform = AddMetaPaths(metapaths, weighted=True)
-    metapath_data = transform(copy.copy(data))
-
-    # Make sure manually added metapaths compute the correct number of edges
-    assert metapath_data['a', 'c'].edge_weight.tolist() == [1, 2]
-    assert metapath_data['a', 'd'].edge_weight.tolist() == [1, 2]
-    assert metapath_data['a', 'a'].edge_weight.tolist() == [2, 1, 4, 2]
-
-    # Compute intra-table metapaths efficiently
+    out = transform(copy.copy(data))
+
+    # Make sure manually added metapaths compute the correct number of edges:
+    edge_index = out['a', 'a'].edge_index
+    edge_weight = out['a', 'a'].edge_weight
+    edge_index, edge_weight = coalesce(edge_index, edge_weight)
+    assert edge_index.tolist() == [[0, 0, 1, 1], [0, 1, 0, 1]]
+    assert edge_weight.tolist() == [1, 2, 2, 4]
+
+    edge_index = out['a', 'c'].edge_index
+    edge_weight = out['a', 'c'].edge_weight
+    edge_index, edge_weight = coalesce(edge_index, edge_weight)
+    assert edge_index.tolist() == [[0, 1], [0, 1]]
+    assert edge_weight.tolist() == [1, 2]
+
+    edge_index = out['a', 'd'].edge_index
+    edge_weight = out['a', 'd'].edge_weight
+    edge_index, edge_weight = coalesce(edge_index, edge_weight)
+    assert edge_index.tolist() == [[0, 1], [0, 0]]
+    assert edge_weight.tolist() == [1, 2]
+
+    # Compute intra-table metapaths efficiently:
     metapaths = [[('a', 'b'), ('b', 'c'), ('c', 'd')]]
-    metapath_data = AddMetaPaths(metapaths, weighted=True)(copy.copy(data))
-    metapath_data['d',
-                  'a'].edge_index = metapath_data['a',
-                                                  'd'].edge_index.flip([0])
-    metapath_data['d', 'a'].edge_weight = metapath_data['a', 'd'].edge_weight
+    out = AddMetaPaths(metapaths, weighted=True)(copy.copy(data))
+    out['d', 'a'].edge_index = out['a', 'd'].edge_index.flip([0])
+    out['d', 'a'].edge_weight = out['a', 'd'].edge_weight
     metapaths = [[('a', 'd'), ('d', 'a')]]
-    metapath_data = AddMetaPaths(metapaths, weighted=True)(metapath_data)
-    del metapath_data['a', 'd']
-    del metapath_data['d', 'a']
-    assert metapath_data['a', 'a'].edge_weight.tolist() == [2, 1, 4, 2]
+    out = AddMetaPaths(metapaths, weighted=True)(out)
+
+    edge_index = out['a', 'a'].edge_index
+    edge_weight = out['a', 'a'].edge_weight
+    edge_index, edge_weight = coalesce(edge_index, edge_weight)
+    assert edge_index.tolist() == [[0, 0, 1, 1], [0, 1, 0, 1]]
+    assert edge_weight.tolist() == [1, 2, 2, 4]
 
 
 def test_add_random_metapaths():