Needed for save and load of gemma 1b (#1903)

stack-info: PR: #1903, branch: drisspg/stack/43

Author: drisspg

test/dtypes/test_affine_quantized.py

@@ -7,6 +7,7 @@
 import unittest
 
 import torch
+import torch.nn as nn
 from torch.testing._internal import common_utils
 from torch.testing._internal.common_utils import (
     TestCase,
@@ -16,6 +17,7 @@
 from torchao.core.config import AOBaseConfig
 from torchao.dtypes import CutlassInt4PackedLayout, Int4CPULayout, SemiSparseLayout
 from torchao.quantization import (
+    Int8DynamicActivationInt8WeightConfig,
     float8_weight_only,
     int4_dynamic_activation_int4_weight,
     int4_weight_only,
@@ -298,6 +300,13 @@ def test_flatten_unflatten(self, device, dtype):
             reconstruct_res = ql(*example_inputs)
             self.assertEqual(reconstruct_res, ref)
 
+    @common_utils.parametrize("device", COMMON_DEVICES)
+    @common_utils.parametrize("dtype", COMMON_DTYPES)
+    def test_alias(self, device, dtype):
+        dummy = nn.Linear(128, 256, dtype=dtype, device=device)
+        quantize_(dummy, Int8DynamicActivationInt8WeightConfig())
+        _ = dummy.weight[...]
+
 
 common_utils.instantiate_parametrized_tests(TestAffineQuantized)
 common_utils.instantiate_parametrized_tests(TestAffineQuantizedBasic)

torchao/dtypes/affine_quantized_tensor_ops.py

@@ -341,7 +341,7 @@ def _(func, types, args, kwargs):
         return func(input_tensor, weight_tensor)
 
 
-@implements(aten.detach.default)
+@implements([aten.detach.default, aten.alias.default])
 def _(func, types, args, kwargs):
     return return_and_correct_aliasing(
         func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)

torchao/quantization/linear_activation_quantized_tensor.py

@@ -184,10 +184,10 @@ def _(func, types, args, kwargs):
         return func(qtensor, original_weight_tensor)
 
 
-@implements(aten.detach.default)
+@implements([aten.detach.default, aten.alias.default])
 def _(func, types, args, kwargs):
     return return_and_correct_aliasing(
-        func, args, kwargs, args[0]._apply_fn_to_data(torch.detach)
+        func, args, kwargs, args[0]._apply_fn_to_data(func)
     )