comments

Author: qihqi

torchax/test/test_interop.py

@@ -4,27 +4,6 @@
 from torchax import interop
 import torchax
 
-class M1(torch.nn.Module):
-
-    def __init__(self):
-        super().__init__()
-        self.x = torch.ones(10, 10)
-
-class M(torch.nn.Module):
-
-    def __init__(self):
-        super().__init__()
-        self.a = torch.nn.Linear(100, 100)
-        self.b = torch.nn.Parameter(
-            torch.ones(10, 10)
-        )
-        c = torch.ones(10, 10)
-        self.register_buffer('c', c)
-        self.register_buffer('c2', c, persistent=False)
-        self.d = torch.ones(10, 10)
-        self.m1 = M1()
-
-
 
 class InteropTest(unittest.TestCase):
 
@@ -33,7 +12,28 @@ def setUp(self):
 
 
     def test_mod_attr(self):
-        m = M()
+
+        class Child(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.x = torch.ones(10, 10)
+
+        class ModWithUnregisteredTensor(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.a = torch.nn.Linear(100, 100)
+                self.b = torch.nn.Parameter(
+                    torch.ones(10, 10)
+                )
+                c = torch.ones(10, 10)
+                self.register_buffer('c', c)
+                self.register_buffer('c2', c, persistent=False)
+                self.d = torch.ones(10, 10)
+                self.m1 = Child()
+
+        m = ModWithUnregisteredTensor()
         params, buffers = interop.extract_all_buffers(m)
         self.assertEqual(
             set(params.keys()), {'a.weight', 'a.bias', 'b'}
@@ -49,7 +49,8 @@ def test_mod_attr(self):
 
     def test_module_with_shared_weights(self):
 
-        class M2(torch.nn.Module):
+        # arrange
+        class Module(torch.nn.Module):
 
             def __init__(self):
                 super().__init__()
@@ -59,23 +60,25 @@ def __init__(self):
             def forward(self, x):
                 return self.a(self.b(x))
 
-        m = M2().to('jax')
+        m = Module().to('jax')
 
         m_jitted = interop.JittableModule(m, dedup_parameters=True)
 
-
         # a's weights and bias and b's weights and bias
         self.assertEqual(len(m.state_dict()), 4)
 
         # b's weights and bias are deduped
         self.assertEqual(len(m_jitted.params), 2)
-
         x = torch.randn(10, 10).to('jax')
-
         expected = m(x)
+
+        # act 
+        actual = m_jitted(x)
 
-        torch.testing.assert_allclose(m_jitted(x), expected)
+        # assert
+        torch.testing.assert_allclose(actual, expected)
 
+        # arrange
         # make sure buffer donation works
         functional_forward = interop.jax_jit(
             functools.partial(m_jitted.functional_call, 'forward'),
@@ -84,8 +87,10 @@ def forward(self, x):
             }
         )
 
-        torch.testing.assert_allclose(
-            functional_forward(m_jitted.params, m_jitted.buffers, x) , expected)
+        # act
+        actual = functional_forward(m_jitted.params, m_jitted.buffers, x)
+        # assert
+        torch.testing.assert_allclose(actual, expected)