Merge branch 'split_implements' of https://github.com/Krishn1412/ao into split_implements

Author: Krishn1412

test/test_utils.py

@@ -389,6 +389,54 @@ def fake_linear(f, types, args, kwargs):
             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")
 
+        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__":
     unittest.main()