Repository: pytorch/PiPPy

Issue:
Implement kwargs support on Pipe.forward()
Currently we don't accept kwargs in `Pipe.forward`: https://github.com/jamesr66a/PiPPy/blob/527af1fd8123d35bd81b9fe304a8d0ed29c9fd8d/pippy/IR.py#L229

We should allow this by looking at the `placeholder` nodes on `Pipe.split_gm` and forming a schema from this, then binding kwarg values to the positional slots for the schema. (TODO: can this be upstreamed to torch.fx?)


Discussion:

Resolving patch:
diff --git a/pippy/IR.py b/pippy/IR.py
--- a/pippy/IR.py
+++ b/pippy/IR.py
@@ -218,10 +218,21 @@ def throw(self, *args, **kwargs):
         self.split_gm.forward = throw
 
     def forward(self, *args, **kwargs):
+        executor_args = args
         if len(kwargs) > 0:
-            # TODO
-            raise NotImplementedError('Kwargs not implemented')
-        return self.executor.run(*args)
+            from inspect import Signature, Parameter
+            parameters = []
+            for node in self.split_gm.graph.nodes:
+                if node.op == 'placeholder':
+                    if node.args and len(node.args) > 0:
+                        parameters.append(Parameter(node.target, Parameter.POSITIONAL_OR_KEYWORD, default=node.args[0]))
+                    else:
+                        parameters.append(Parameter(node.target, Parameter.POSITIONAL_OR_KEYWORD))
+            signature = Signature(parameters)
+            ba = signature.bind(*args, **kwargs)
+            ba.apply_defaults()
+            executor_args = ba.arguments.values()
+        return self.executor.run(*executor_args)
 
     def remap_qualname(self, qualname):
         # TODO: annoying


Test patch:
diff --git a/test/test_ir.py b/test/test_ir.py
--- a/test/test_ir.py
+++ b/test/test_ir.py
@@ -327,5 +327,42 @@ def forward(self, x):
         with self.assertRaisesRegex(AttributeError, 'Specified target foo.bar.baz referenced nonexistent module foo'):
             annotate_split_points(m, {'foo.bar.baz': PipeSplitWrapper.SplitPoint.END})
 
+    def test_pipe_forward_kwargs(self):
+        class ModuleWithArgsAndKWArgs(torch.nn.Module):
+            def forward(self, x, y=2):
+                return x + y
+
+        module = ModuleWithArgsAndKWArgs()
+        pipe = Pipe.from_tracing(module)
+        self.assertEqual(pipe(1), 3)
+        self.assertEqual(pipe(1, y=3), 4)
+        self.assertEqual(pipe(1, 4), 5)
+        with self.assertRaisesRegex(TypeError, "got an unexpected keyword argument 'z'"):
+            pipe(1, z=100)
+        with self.assertRaisesRegex(TypeError, "got an unexpected keyword argument 'a'"):
+            pipe(1, y=3, a='a', b=True, c='c')
+        with self.assertRaisesRegex(TypeError, "multiple values for argument 'y'"):
+            pipe(1, 4, y=100)
+
+        class ModuleWithOnlyKWArgs(torch.nn.Module):
+            def forward(self, x=1, y=2):
+                return x + y
+
+        module = ModuleWithOnlyKWArgs()
+        pipe = Pipe.from_tracing(module)
+        self.assertEqual(pipe(), 3)
+        self.assertEqual(pipe(2), 4)
+        self.assertEqual(pipe(x=3), 5)
+        self.assertEqual(pipe(y=4), 5)
+        self.assertEqual(pipe(3, 3), 6)
+        self.assertEqual(pipe(3, y=4), 7)
+        self.assertEqual(pipe(x=4, y=4), 8)
+        with self.assertRaisesRegex(TypeError, "got an unexpected keyword argument 'b'"):
+            pipe(b=True, a='a', c='c')
+        with self.assertRaisesRegex(TypeError, "multiple values for argument 'x'"):
+            pipe(2, x=3)
+        with self.assertRaisesRegex(TypeError, "multiple values for argument 'x'"):
+            pipe(1, 2, x=3, y=4)
+
 if __name__ == '__main__':
     unittest.main()