Add test_make_fx_model_train example (#82011)

Summary:
X-link: pytorch/pytorch#82011

Pull Request resolved: #980

Test Plan: CI should pass

Reviewed By: benoitsteiner

Differential Revision: D38078694

Pulled By: mostafaelhoushi

fbshipit-source-id: bce8ccebce3b604dc5b98168b783aa6e21938f95

Author: mostafaelhoushi

test/test_pythonkey.py

@@ -11,6 +11,8 @@
 import unittest
 import warnings
 import itertools
+import torch.nn.utils._stateless as stateless
+from collections.abc import Iterable
 from functools import partial
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from functorch import (
@@ -74,6 +76,83 @@ def f(x):
         new_inp = torch.randn(3)
         self.assertEqual(fx_f(new_inp), f(new_inp))
 
+    def test_make_fx_model_fwd_bwd(self, device):
+        class Foo(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(5, 5)
+
+            def forward(self, x):
+                return self.linear(x).relu()
+
+        model = Foo()
+
+        def f(x, params):
+            out = stateless.functional_call(model, params, x).sum()
+            out.backward()
+            return list(params.values())
+        input = torch.randn(3, 5, requires_grad=True)
+        params = dict(model.named_parameters())
+        fx_f = make_fx(f)(input, params)
+        # TODO: what assert statement should we add here?
+        assert(fx_f(input, params) is not None)
+
+    def test_make_fx_model_fwd_bwd_wgtupdate(self, device):
+        class Foo(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(5, 5)
+
+            def forward(self, x):
+                return self.linear(x).relu()
+
+        model = Foo()
+
+        def f(args, params, buffers):
+            if not isinstance(args, Iterable):
+                args = [args]
+            params_and_buffers = {**params, **buffers}
+            out = stateless.functional_call(model, params_and_buffers, args)
+            out.sum().backward()
+            return [p - 1e-4 * p.grad for p in params.values()]
+
+        input = torch.randn(3, 5, requires_grad=True)
+        params = dict(model.named_parameters())
+        buffers = dict(model.named_buffers())
+        fx_f = make_fx(f)(input, params, buffers)
+        # TODO: what assert statement should we add here?
+        assert(fx_f(input, params, buffers) is not None)
+
+    def test_make_fx_model_train_with_optim(self, device):
+        class Foo(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(5, 5)
+
+            def forward(self, x):
+                return self.linear(x).relu()
+
+        model = Foo()
+        optim = torch.optim.SGD(model.parameters(), lr=1e-4)
+
+        def f(args, params, buffers):
+            if not isinstance(args, Iterable):
+                args = [args]
+            params_and_buffers = {**params, **buffers}
+            out = stateless.functional_call(model, params_and_buffers, args)
+            out.sum().backward()
+            optim.step()
+
+            # TODO: this causes graph to show an output with many incoming edges. Shall we try `return None` or simply don't return?
+            return list(params.values())
+
+        input = torch.randn(3, 5, requires_grad=True)
+        params = dict(model.named_parameters())
+        buffers = dict(model.named_buffers())
+        fx_f = make_fx(f)(input, params, buffers)
+        # TODO: what assert statement should we add here?
+        assert(fx_f(input, params, buffers) is not None)
+
     def test_scalar_device(self, device):
         def f(a, b):
             return a + b