test: add test for empty parameters

XuehaiPan · XuehaiPan · commit fd1642c478ef · 2023-03-04T00:54:46.000+08:00
diff --git a/tests/test_implicit.py b/tests/test_implicit.py
@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import copy
+import re
 from collections import OrderedDict
 from types import FunctionType
 
@@ -698,8 +699,8 @@ def __init__(self, x):
             super().__init__()
             self.x = x
 
-        def optimality(self):
-            return (self.x,)
+        def objective(self):
+            return self.x
 
         def solve(self):
             pass
@@ -717,6 +718,10 @@ def solve(self):
     with pytest.raises(RuntimeError, match='The module has no parameters.'):
         model.solve()
 
+    model = EmptyParameters(torch.zeros(8, requires_grad=True))
+    with pytest.raises(RuntimeError, match='The module has no parameters.'):
+        model.optimality()
+
     model = EmptyParameters(torch.zeros(8, requires_grad=True))
     model.register_parameter('y', torch.zeros(8, requires_grad=True))
     model.solve()
@@ -728,3 +733,124 @@ def solve(self):
     model = EmptyParameters(nn.Linear(8, 8))
     model.register_parameter('y', torch.zeros(8, requires_grad=True))
     model.solve()
+
+
+def test_module_enable_implicit_gradients_twice() -> None:
+    class MyModule(torchopt.nn.ImplicitMetaGradientModule):
+        def objective(self):
+            return torch.tensor(0.0)
+
+        def solve(self):
+            pass
+
+    from torchopt.diff.implicit.nn.module import enable_implicit_gradients
+
+    with pytest.raises(
+        TypeError,
+        match='Implicit gradients are already enabled for the `solve` method.',
+    ):
+        enable_implicit_gradients(MyModule)
+
+
+def test_module_abstract_methods() -> None:
+    class MyModule1(torchopt.nn.ImplicitMetaGradientModule):
+        def objective(self):
+            return torch.tensor(0.0)
+
+    with pytest.raises(
+        TypeError,
+        match="Can't instantiate abstract class",
+    ):
+        MyModule1()
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape(
+            'ImplicitMetaGradientModule requires either an optimality() method or an objective() method'
+        ),
+    ):
+
+        class MyModule2(torchopt.nn.ImplicitMetaGradientModule):
+            def solve(self):
+                pass
+
+    class MyModule3(torchopt.nn.ImplicitMetaGradientModule):
+        def optimality(self):
+            return ()
+
+        def solve(self):
+            pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape('method optimality() must not be a staticmethod.'),
+    ):
+
+        class MyModule4(torchopt.nn.ImplicitMetaGradientModule):
+            @staticmethod
+            def optimality():
+                return ()
+
+            def solve(self):
+                pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape('method optimality() must not be a classmethod.'),
+    ):
+
+        class MyModule5(torchopt.nn.ImplicitMetaGradientModule):
+            @classmethod
+            def optimality(self):
+                return ()
+
+            def solve(self):
+                pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape('method optimality() must be callable.'),
+    ):
+
+        class MyModule6(torchopt.nn.ImplicitMetaGradientModule):
+            optimality = 0
+
+            def solve(self):
+                pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape('method optimality() must not be a staticmethod.'),
+    ):
+
+        class MyModule7(torchopt.nn.ImplicitMetaGradientModule):
+            @staticmethod
+            def optimality():
+                return ()
+
+            def solve(self):
+                pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape('method objective() must not be a classmethod.'),
+    ):
+
+        class MyModule8(torchopt.nn.ImplicitMetaGradientModule):
+            @classmethod
+            def objective(self):
+                return ()
+
+            def solve(self):
+                pass
+
+    with pytest.raises(
+        TypeError,
+        match=re.escape('method objective() must be callable.'),
+    ):
+
+        class MyModule9(torchopt.nn.ImplicitMetaGradientModule):
+            objective = 0
+
+            def solve(self):
+                pass
diff --git a/tests/test_zero_order.py b/tests/test_zero_order.py
@@ -14,6 +14,7 @@
 # ==============================================================================
 
 import functorch
+import pytest
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -117,3 +118,55 @@ def sample(self, sample_shape=torch.Size()):
         optimizer.zero_grad()
         loss.backward()  # compute gradients
         optimizer.step()  # update network parameters
+
+
+def test_module_enable_zero_order_gradients_twice() -> None:
+    class MyModule(torchopt.nn.ZeroOrderGradientModule):
+        def forward(self):
+            return torch.tensor(0.0)
+
+        def sample(self, sample_shape):
+            return torch.tensor(0.0)
+
+    from torchopt.diff.zero_order.nn.module import enable_zero_order_gradients
+
+    with pytest.raises(
+        TypeError,
+        match='Zero-order gradient estimation is already enabled for the `forward` method.',
+    ):
+        enable_zero_order_gradients(MyModule)
+
+
+def test_module_empty_parameters() -> None:
+    class MyModule(torchopt.nn.ZeroOrderGradientModule):
+        def forward(self):
+            return torch.tensor(0.0)
+
+        def sample(self, sample_shape):
+            return torch.tensor(0.0)
+
+    m = MyModule()
+    with pytest.raises(RuntimeError, match='The module has no parameters.'):
+        m()
+
+
+def test_module_abstract_methods() -> None:
+    class MyModule1(torchopt.nn.ZeroOrderGradientModule):
+        def forward(self):
+            return torch.tensor(0.0)
+
+    with pytest.raises(
+        TypeError,
+        match="Can't instantiate abstract class",
+    ):
+        MyModule1()
+
+    class MyModule2(torchopt.nn.ZeroOrderGradientModule):
+        def sample(self, sample_shape):
+            return torch.tensor(0.0)
+
+    with pytest.raises(
+        TypeError,
+        match="Can't instantiate abstract class",
+    ):
+        MyModule2()
diff --git a/torchopt/diff/implicit/nn/module.py b/torchopt/diff/implicit/nn/module.py
@@ -20,6 +20,7 @@
 
 import abc
 import functools
+import inspect
 import itertools
 from typing import Any, Iterable
 
@@ -79,10 +80,9 @@ def make_optimality_from_objective(
     cls: type[ImplicitMetaGradientModule],
 ) -> type[ImplicitMetaGradientModule]:
     """Derives the optimality function of the objective function."""
-    if (
-        getattr(cls, 'objective', ImplicitMetaGradientModule.objective)
-        is ImplicitMetaGradientModule.objective
-    ):
+    static_super_objective = inspect.getattr_static(ImplicitMetaGradientModule, 'objective')
+    static_cls_optimality = inspect.getattr_static(cls, 'optimality', static_super_objective)
+    if static_cls_optimality is static_super_objective:
         raise TypeError('The objective function is not defined.')
 
     def optimality(self: ImplicitMetaGradientModule, *input: Any, **kwargs: Any) -> TupleOfTensors:
@@ -167,7 +167,7 @@ def wrapped(self: ImplicitMetaGradientModule, *input: Any, **kwargs: Any) -> Any
     return cls
 
 
-class ImplicitMetaGradientModule(MetaGradientModule):
+class ImplicitMetaGradientModule(MetaGradientModule, metaclass=abc.ABCMeta):
     """The base class for differentiable implicit meta-gradient models."""
 
     _custom_optimality: bool
@@ -179,28 +179,30 @@ def __init_subclass__(cls, linear_solve: LinearSolver | None = None) -> None:
         super().__init_subclass__()
         cls.linear_solve = linear_solve
 
-        optimality = getattr(cls, 'optimality', ImplicitMetaGradientModule.optimality)
-        objective = getattr(cls, 'objective', ImplicitMetaGradientModule.objective)
-        cls._custom_optimality = optimality is not ImplicitMetaGradientModule.optimality
-        cls._custom_objective = objective is not ImplicitMetaGradientModule.objective
+        static_super_optimality = inspect.getattr_static(ImplicitMetaGradientModule, 'optimality')
+        static_super_objective = inspect.getattr_static(ImplicitMetaGradientModule, 'objective')
+        static_cls_optimality = inspect.getattr_static(cls, 'optimality')
+        static_cls_objective = inspect.getattr_static(cls, 'objective')
+        cls._custom_optimality = static_cls_optimality is not static_super_optimality
+        cls._custom_objective = static_cls_objective is not static_super_objective
 
         if cls._custom_optimality:
-            if isinstance(optimality, staticmethod):
+            if isinstance(static_cls_optimality, staticmethod):
                 raise TypeError('method optimality() must not be a staticmethod.')
-            if isinstance(optimality, classmethod):
+            if isinstance(static_cls_optimality, classmethod):
                 raise TypeError('method optimality() must not be a classmethod.')
-            if not callable(optimality):
+            if not callable(static_cls_optimality):
                 raise TypeError('method optimality() must be callable.')
         elif not cls._custom_objective:
             raise TypeError(
                 'ImplicitMetaGradientModule requires either an optimality() method or an objective() method'
             )
         else:
-            if isinstance(objective, staticmethod):
+            if isinstance(static_cls_objective, staticmethod):
                 raise TypeError('method objective() must not be a staticmethod.')
-            if isinstance(objective, classmethod):
+            if isinstance(static_cls_objective, classmethod):
                 raise TypeError('method objective() must not be a classmethod.')
-            if not callable(objective):
+            if not callable(static_cls_objective):
                 raise TypeError('method objective() must be callable.')
 
             make_optimality_from_objective(cls)
