[API Compatibility] Add Four api: torch.nn.Module.register_forward_hook, torch.nn.Module.register_forward_pre_hook, torch.nn.modules.module._IncompatibleKeys and torch.optim.lr_scheduler.LRScheduler

paconvert/api_mapping.json

@@ -6772,36 +6772,10 @@
     "Matcher": "ChangePrefixMatcher"
   },
   "torch.nn.Module.register_forward_hook": {
-    "Matcher": "GenericMatcher",
-    "paddle_api": "paddle.nn.Module.register_forward_post_hook",
-    "args_list": [
-      "hook",
-      "*",
-      "prepend",
-      "with_kwargs",
-      "always_call"
-    ],
-    "unsupport_args": [
-      "prepend",
-      "with_kwargs",
-      "always_call"
-    ],
-    "min_input_args": 1
+    "Matcher": "ChangePrefixMatcher"
   },
   "torch.nn.Module.register_forward_pre_hook": {
-    "Matcher": "GenericMatcher",
-    "paddle_api": "paddle.nn.Module.register_forward_pre_hook",
-    "args_list": [
-      "hook",
-      "*",
-      "prepend",
-      "with_kwargs"
-    ],
-    "unsupport_args": [
-      "prepend",
-      "with_kwargs"
-    ],
-    "min_input_args": 1
+    "Matcher": "ChangePrefixMatcher"
   },
   "torch.nn.Module.register_full_backward_hook": {
     "min_input_args": 1
@@ -8192,6 +8166,9 @@
   "torch.nn.modules.module.Module": {
     "Matcher": "ChangePrefixMatcher"
   },
+  "torch.nn.modules.module._IncompatibleKeys": {
+    "Matcher": "ChangePrefixMatcher"
+  },
   "torch.nn.modules.utils._ntuple": {
     "Matcher": "NTupleMatcher",
     "args_list": [
@@ -8773,6 +8750,9 @@
   "torch.optim.lr_scheduler.ExponentialLR": {
     "Matcher": "ChangePrefixMatcher"
   },
+  "torch.optim.lr_scheduler.LRScheduler": {
+    "Matcher": "ChangePrefixMatcher"
+  },
   "torch.optim.lr_scheduler.LambdaLR": {
     "Matcher": "ChangePrefixMatcher"
   },

tests/test_nn_Module_register_forward_hook.py

@@ -102,11 +102,7 @@ def hook(module, fea_in, fea_out):
         net(a)
         """
     )
-    obj.run(
-        pytorch_code,
-        unsupport=True,
-        reason="prepend, with_kwargs and always_call is not supported",
-    )
+    obj.run(pytorch_code, ["result"])
 
 
 def test_case_5():
@@ -132,8 +128,173 @@ def hook(module, fea_in, fea_out):
         net(a)
         """
     )
+    obj.run(pytorch_code, ["result"])
+
+
+def test_case_6():
+    """prepend=True runs the new hook before existing hooks."""
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        events = []
+
+        class DoubleNum(torch.nn.Module):
+            def forward(self, x):
+                return 2 * x
+
+        def first_hook(module, args, output):
+            events.append("first")
+            return output + 1
+
+        def second_hook(module, args, output):
+            events.append("second")
+            return output * 2
+
+        net = DoubleNum()
+        net.register_forward_hook(second_hook)
+        net.register_forward_hook(first_hook, prepend=True)
+        x = torch.tensor(
+            [[-1.5, 2.0], [0.25, -3.0]],
+            dtype=torch.float64,
+        )
+        result = net(x)
+        """
+    )
+    obj.run(pytorch_code, ["events", "result"])
+
+
+def test_case_7():
+    """with_kwargs=True passes forward keyword arguments to the hook."""
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        observed_kwargs = []
+
+        class ScaleAndShift(torch.nn.Module):
+            def forward(self, x, scale=1.0, offset=0.0):
+                return x * scale + offset
+
+        def hook(module, args, kwargs, output):
+            observed_kwargs.append((kwargs["scale"], kwargs["offset"]))
+            return output + kwargs["scale"]
+
+        net = ScaleAndShift()
+        net.register_forward_hook(hook, with_kwargs=True)
+        x = torch.tensor(
+            [[[-1.0, 0.5], [2.0, -3.0]]],
+            dtype=torch.float32,
+        )
+        result = net(x, scale=2.5, offset=-1.0)
+        """
+    )
+    obj.run(pytorch_code, ["observed_kwargs", "result"])
+
+
+def test_case_8():
+    """always_call=True invokes the hook when forward raises."""
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        events = []
+
+        class BrokenModule(torch.nn.Module):
+            def forward(self, x):
+                raise RuntimeError("forward failed")
+
+        def hook(module, args, output):
+            events.append(output is None)
+
+        net = BrokenModule()
+        net.register_forward_hook(hook, always_call=True)
+        try:
+            net(torch.tensor([1.0, -2.0]))
+        except RuntimeError:
+            pass
+        """
+    )
+    obj.run(pytorch_code, ["events"])
+
+
+def test_case_9():
+    """The returned handle removes the registered hook."""
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        events = []
+
+        class AddOne(torch.nn.Module):
+            def forward(self, x):
+                return x + 1
+
+        def hook(module, args, output):
+            events.append("called")
+            return output * 2
+
+        net = AddOne()
+        handle = net.register_forward_hook(hook)
+        result_before_remove = net(torch.tensor([1, -2, 3]))
+        handle.remove()
+        result_after_remove = net(torch.tensor([1, -2, 3]))
+        """
+    )
     obj.run(
         pytorch_code,
-        unsupport=True,
-        reason="prepend, with_kwargs and always_call is not supported",
+        ["events", "result_before_remove", "result_after_remove"],
+    )
+
+
+def test_case_10():
+    """The hook can be supplied through variable positional arguments."""
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        class DoubleNum(torch.nn.Module):
+            def forward(self, x):
+                return 2 * x
+
+        def hook(module, args, output):
+            return output - 1
+
+        net = DoubleNum()
+        hook_args = (hook,)
+        net.register_forward_hook(*hook_args)
+        x = torch.tensor([[[-2.0, 1.0], [0.5, -0.25]]])
+        result = net(x)
+        """
+    )
+    obj.run(pytorch_code, ["result"])
+
+
+def test_case_11():
+    """All arguments can be supplied through variable keyword arguments."""
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        events = []
+
+        class DoubleNum(torch.nn.Module):
+            def forward(self, x):
+                return 2 * x
+
+        def hook(module, args, kwargs, output):
+            events.append(len(kwargs))
+            return output + 3
+
+        net = DoubleNum()
+        hook_kwargs = {
+            "hook": hook,
+            "prepend": False,
+            "with_kwargs": True,
+            "always_call": False,
+        }
+        net.register_forward_hook(**hook_kwargs)
+        result = net(torch.tensor([-1.0, 2.0]))
+        """
     )
+    obj.run(pytorch_code, ["events", "result"])