Testing subclassing

test/test_utils.py

@@ -392,6 +392,49 @@ def fake_linear(func, types, args, kwargs):
             counter["calls"], 2, "Expected fake_linear to be called via aten.t.default"
         )
 
+    def test_subclassing(self):
+        class Parent(TorchAOBaseTensor):
+            tensor_data_names = ["qdata"]
+            tensor_attribute_names = ["attr"]
+
+        Parent._ATEN_OP_TABLE[Parent]["op_parent"] = "parent_impl"
+        Parent._TORCH_FN_TABLE[Parent]["fn_parent"] = "parent_fn_impl"
+
+        class Child(Parent):
+            tensor_data_names = ["qdata"]
+            tensor_attribute_names = ["attr"]
+
+        # ensure child has copied parent ops
+        self.assertEqual(Child._ATEN_OP_TABLE[Child]["op_parent"], "parent_impl")
+        self.assertEqual(Child._TORCH_FN_TABLE[Child]["fn_parent"], "parent_fn_impl")
+
+        # ensure the top-level dicts are distinct (not inherited)
+        self.assertIsNot(Parent._ATEN_OP_TABLE, Child._ATEN_OP_TABLE)
+        self.assertIsNot(Parent._TORCH_FN_TABLE, Child._TORCH_FN_TABLE)
+
+        # change the parent's op after subclass creation — should not leak
+        Parent._ATEN_OP_TABLE[Parent]["new_op"] = "added_later"
+        self.assertNotIn("new_op", Child._ATEN_OP_TABLE[Child])
+
+    def test_multiple_inheritance(self):
+        class A(TorchAOBaseTensor):
+            tensor_data_names = ["a"]
+            tensor_attribute_names = ["b"]
+
+        class B(TorchAOBaseTensor):
+            tensor_data_names = ["a"]
+            tensor_attribute_names = ["b"]
+
+        A._ATEN_OP_TABLE[A]["shared"] = "from_a"
+        B._ATEN_OP_TABLE[B]["shared"] = "from_b"
+
+        class C(A, B):
+            tensor_data_names = ["a"]
+            tensor_attribute_names = ["b"]
+
+        # C(A, B) should inherit from A then B, so B wins
+        self.assertEqual(C._ATEN_OP_TABLE[C]["shared"], "from_b")
+
 
 if __name__ == "__main__":
     unittest.main()