Add 2 ops

Author: qihqi

torchax/test/test_interop.py

@@ -1,6 +1,8 @@
+import functools
 import torch
 import unittest
 from torchax import interop
+import torchax
 
 class M1(torch.nn.Module):
 
@@ -23,8 +25,12 @@ def __init__(self):
         self.m1 = M1()
 
 
+
 class InteropTest(unittest.TestCase):
 
+    def setUp(self):
+        torchax.enable_globally()
+
 
     def test_mod_attr(self):
         m = M()
@@ -41,6 +47,49 @@ def test_mod_attr(self):
         self.assertEqual(m.a.weight.item(), 0)
         self.assertEqual(m.m1.x.item(), 0)
 
+    def test_module_with_shared_weights(self):
+
+        class M2(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.a = torch.nn.Linear(10, 10)
+                self.b = self.a
+            
+            def forward(self, x):
+                return self.a(self.b(x))
+
+        m = M2().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)
+        
+        torch.testing.assert_allclose(m_jitted(x), expected)
+
+        # make sure buffer donation works
+        functional_forward = interop.jax_jit(
+            functools.partial(m_jitted.functional_call, 'forward'),
+            kwargs_for_jax_jit={
+                'donate_argnums': (0, )
+            }
+        )
+
+        torch.testing.assert_allclose(
+            functional_forward(m_jitted.params, m_jitted.buffers, x) , expected)
+
+
+        
+
 
 
 if __name__ == '__main__':

torchax/torchax/interop.py

@@ -50,7 +50,7 @@ def set_one(module, prefix):
 
 class JittableModule(torch.nn.Module):
 
-    def __init__(self, m: torch.nn.Module, extra_jit_args={}, dedup_paramters=True):
+    def __init__(self, m: torch.nn.Module, extra_jit_args={}, dedup_parameters=True):
         super().__init__()
         self.params, self.buffers = extract_all_buffers(m)
         self._model = m
@@ -60,7 +60,7 @@ def __init__(self, m: torch.nn.Module, extra_jit_args={}, dedup_paramters=True):
 
         self._extra_dumped_weights = {}
 
-        if dedup_paramters:
+        if dedup_parameters:
             temp = collections.defaultdict(list)
             for k, v in self.params.items():
                 temp[id(v)].append(k)
@@ -72,9 +72,6 @@ def __init__(self, m: torch.nn.Module, extra_jit_args={}, dedup_paramters=True):
                     for extra_keys in v[1:]:
                         del self.params[extra_keys]
 
-
-
-
     def __call__(self, *args, **kwargs):
         return self.forward(*args, **kwargs)
 

torchax/torchax/ops/jaten.py

@@ -1590,12 +1590,14 @@ def _aten_bitwise_not(self):
 
 
 # aten.bitwise_left_shift
+@op(torch.ops.aten.__lshift__)
 @op(torch.ops.aten.bitwise_left_shift)
 def _aten_bitwise_left_shift(input, other):
   return jnp.left_shift(input, other)
 
 
 # aten.bitwise_right_shift
+@op(torch.ops.aten.__rshift__)
 @op(torch.ops.aten.bitwise_right_shift)
 def _aten_bitwise_right_shift(input, other):
   return jnp.right_shift(input, other)