Updating test to check the user code

Author: Krishn1412

test/test_utils.py

@@ -121,47 +121,53 @@ def __init__(self, qdata, attr, device=None):
         # after copy_, the tensor values should match
         self.assertEqual(lp_tensor.qdata[0], another_lp_tensor.qdata[0])
 
-    def test_implements_and_torch_function_together(self):
-        """Ensure a function decorated with both @_implements and @_implements_torch_function works."""
-
-        implements = TorchAOBaseTensor.implements
-        implements_torch_function = TorchAOBaseTensor.implements_torch_function
-
-        @implements([torch.ops.aten.linear.default])
-        @implements_torch_function([F.linear])
-        def fake_linear(f, types, args, kwargs):
-            x, w, b = args
-            return torch.matmul(x, w.T) + (b if b is not None else 0)
-
-        # make sure both got registered on TorchAOBaseTensor
-        self.assertIn(
-            torch.ops.aten.linear.default,
-            TorchAOBaseTensor._ATEN_OP_TABLE[TorchAOBaseTensor],
-        )
-        self.assertIn(F.linear, TorchAOBaseTensor._TORCH_FN_TABLE[TorchAOBaseTensor])
-
-        # check they both point to the same function wrapper
-        aten_wrapper = TorchAOBaseTensor._ATEN_OP_TABLE[TorchAOBaseTensor][
-            torch.ops.aten.linear.default
-        ]
-        torchfn_wrapper = TorchAOBaseTensor._TORCH_FN_TABLE[TorchAOBaseTensor][F.linear]
-
-        # check if they wrap the same underlying function
-        self.assertEqual(aten_wrapper.__wrapped__, fake_linear)
-        self.assertEqual(torchfn_wrapper.__wrapped__, fake_linear)
-
-        # run through the wrapper
-        x = torch.randn(2, 3)
-        w = torch.randn(4, 3)
-        b = torch.randn(4)
-
-        out_aten = aten_wrapper(fake_linear, (TorchAOBaseTensor,), (x, w, b), {})
-        out_torchfn = torchfn_wrapper(fake_linear, (TorchAOBaseTensor,), (x, w, b), {})
-
-        expected = F.linear(x, w, b)
-
-        self.assertTrue(torch.allclose(out_aten, expected, atol=1e-6))
-        self.assertTrue(torch.allclose(out_torchfn, expected, atol=1e-6))
+        def test_implements_and_torch_function_together(self):
+            """Ensure a function decorated with both @_implements and @_implements_torch_function works.
+            """
+            counter = {"calls": 0}
+
+            class MyTensor(TorchAOBaseTensor):
+                tensor_data_names = ["qdata"]
+                tensor_attribute_names = ["attr", "device"]
+
+                def __new__(cls, qdata, attr="attr", device=None):
+                    shape = qdata.shape
+                    if device is None:
+                        device = qdata.device
+                    kwargs = {"device": device}
+                    return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)
+
+                def __init__(self, qdata, attr="attr", device=None):
+                    self.qdata = qdata
+                    self.attr = attr
+
+            # Register the same implementation for both aten and torch-level function
+            @MyTensor.implements(torch.ops.aten.linear.default)
+            @MyTensor.implements_torch_function(F.linear)
+            def fake_linear(f, types, args, kwargs):
+                x, w, b = args
+                w_plain = getattr(w, "qdata", w)
+                b_plain = getattr(b, "qdata", b) if b is not None else None
+                counter["calls"] += 1
+                return torch.matmul(x, w_plain.T) + (b_plain if b_plain is not None else 0)
+
+            
+            l = torch.nn.Linear(2, 3)
+            orig_w = l.weight.detach().clone()
+            orig_b = l.bias.detach().clone() if l.bias is not None else None
+
+            p = torch.nn.Parameter(orig_w)
+            p.data = MyTensor(orig_w, "attr", None)
+            l.weight = p
+
+            x = torch.randn(4, 2)
+
+            # module path (F.linear)
+            self.assertEqual(counter["calls"], 1, "Expected fake_linear to be called once via F.linear")
+
+            # aten path
+            torch.ops.aten.linear.default(x, l.weight, l.bias)
+            self.assertEqual(counter["calls"], 2, "Expected fake_linear to be called once via aten.linear")
 
 
 if __name__ == "__main__":