Repository: pyg-team/pytorch_geometric
Pull request #7209 (2023-04-19 13:11:00)

Added `optimizer_resolver`

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7209?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7209](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7209?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (dde7303) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/fadf0409c9e216e20dbc821d88c8eeea2857bf3e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (fadf040) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head dde7303 differs from pull request most recent head 80514d4. Consider uploading reports for the commit 80514d4 to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #7209   +/-   ##
=======================================
  Coverage   91.58%   91.59%           
=======================================
  Files         436      436           
  Lines       24067    24067           
=======================================
+ Hits        22042    22043    +1     
+ Misses       2025     2024    -1     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7209?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/resolver.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7209?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3Jlc29sdmVyLnB5) | `92.00% <ø> (+2.34%)` | :arrow_up: |
| [torch\_geometric/nn/resolver.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7209?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Jlc29sdmVyLnB5) | `91.17% <100.00%> (+0.55%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7bf92c39e..b35f93d61 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `optimizer_resolver` ([#7209](https://github.com/pyg-team/pytorch_geometric/pull/7209))
 - Added `type_ptr` argument to `HeteroLayerNorm` ([#7208](https://github.com/pyg-team/pytorch_geometric/pull/7208))
 - Added an option to benchmark scripts to write PyTorch profiler results to CSV ([#7114](https://github.com/pyg-team/pytorch_geometric/pull/7114))
 - Added subgraph type sampling option with bidirectional edge support ([#7199](https://github.com/pyg-team/pytorch_geometric/pull/7199), [#7200](https://github.com/pyg-team/pytorch_geometric/pull/7200))
diff --git a/torch_geometric/nn/resolver.py b/torch_geometric/nn/resolver.py
index ce6a94958..099ec3ff7 100644
--- a/torch_geometric/nn/resolver.py
+++ b/torch_geometric/nn/resolver.py
@@ -28,7 +28,6 @@ def swish(x: Tensor) -> Tensor:
 
 
 def activation_resolver(query: Union[Any, str] = 'relu', *args, **kwargs):
-    import torch
     base_cls = torch.nn.Module
     base_cls_repr = 'Act'
     acts = [
@@ -47,8 +46,6 @@ def activation_resolver(query: Union[Any, str] = 'relu', *args, **kwargs):
 
 
 def normalization_resolver(query: Union[Any, str], *args, **kwargs):
-    import torch
-
     import torch_geometric.nn.norm as norm
     base_cls = torch.nn.Module
     base_cls_repr = 'Norm'
@@ -77,6 +74,18 @@ def aggregation_resolver(query: Union[Any, str], *args, **kwargs):
     return resolver(aggrs, aggr_dict, query, base_cls, None, *args, **kwargs)
 
 
+# Optimizer Resolver ##########################################################
+
+
+def optimizer_resolver(query: Union[Any, str], *args, **kwargs):
+    base_cls = Optimizer
+    optimizers = [
+        optimizer for optimizer in vars(torch.optim).values()
+        if isinstance(optimizer, type) and issubclass(optimizer, base_cls)
+    ]
+    return resolver(optimizers, {}, query, base_cls, None, *args, **kwargs)
+
+
 # Learning Rate Scheduler Resolver ############################################
 
 
diff --git a/torch_geometric/resolver.py b/torch_geometric/resolver.py
index 6b642d5bc..d10c36ab4 100644
--- a/torch_geometric/resolver.py
+++ b/torch_geometric/resolver.py
@@ -22,9 +22,7 @@ def resolver(classes: List[Any], class_dict: Dict[str, Any],
         if query_repr == key_repr:
             if inspect.isclass(cls):
                 obj = cls(*args, **kwargs)
-                assert callable(obj)
                 return obj
-            assert callable(cls)
             return cls
 
     for cls in classes:
@@ -32,9 +30,7 @@ def resolver(classes: List[Any], class_dict: Dict[str, Any],
         if query_repr in [cls_repr, cls_repr.replace(base_cls_repr, '')]:
             if inspect.isclass(cls):
                 obj = cls(*args, **kwargs)
-                assert callable(obj)
                 return obj
-            assert callable(cls)
             return cls
 
     choices = set(cls.__name__ for cls in classes) | set(class_dict.keys())


Test patch:
diff --git a/test/nn/test_resolver.py b/test/nn/test_resolver.py
index 46ffd4ed8..cc3dac69e 100644
--- a/test/nn/test_resolver.py
+++ b/test/nn/test_resolver.py
@@ -8,6 +8,7 @@ from torch_geometric.nn.resolver import (
     aggregation_resolver,
     lr_scheduler_resolver,
     normalization_resolver,
+    optimizer_resolver,
 )
 
 
@@ -58,6 +59,22 @@ def test_normalization_resolver(norm_tuple):
                       norm_module)
 
 
+def test_optimizer_resolver():
+    params = [torch.nn.Parameter(torch.Tensor(1))]
+
+    assert isinstance(optimizer_resolver(torch.optim.SGD(params, lr=0.01)),
+                      torch.optim.SGD)
+    assert isinstance(optimizer_resolver(torch.optim.Adam(params)),
+                      torch.optim.Adam)
+    assert isinstance(optimizer_resolver(torch.optim.Rprop(params)),
+                      torch.optim.Rprop)
+
+    assert isinstance(optimizer_resolver('sgd', params, lr=0.01),
+                      torch.optim.SGD)
+    assert isinstance(optimizer_resolver('adam', params), torch.optim.Adam)
+    assert isinstance(optimizer_resolver('rprop', params), torch.optim.Rprop)
+
+
 @pytest.mark.parametrize('scheduler_args', [
     ('constant_with_warmup', LambdaLR),
     ('linear_with_warmup', LambdaLR),