up

Author: metascroy

torchao/experimental/tests/test_embedding_xbit_quantizer.py

@@ -183,10 +183,9 @@ def test_shared_embedding(self):
         self.assertTrue(torch.allclose(result, exported_result))
 
         # Check the shared_embedding and linear ops use the same lifted weight
-        weight = "b_getattr_l__fn_____0___unembedding_packed_weights"
         expected_lines = [
-            f"torch.ops.torchao._shared_embedding_4bit.default({weight}, 4096, 131, 4096, reshape)",
-            f"torch.ops.torchao._linear_8bit_act_4bit_weight.default(linear, {weight}, 4096, 131, 4096)",
+            "torch.ops.torchao._shared_embedding_4bit.default",
+            "torch.ops.torchao._linear_8bit_act_4bit_weight.defaul",
         ]
         for line in expected_lines:
             FileCheck().check_count(line, 1, exactly=True).run(

torchao/quantization/quant_api.py

@@ -2133,6 +2133,7 @@ def _intx_weight_only_quantize_tensor(weight, config):
         zero_point_domain=ZeroPointDomain.INT,
         _layout=layout,
     )
+    return weight
 
 
 @register_quantize_module_handler(IntxWeightOnlyConfig)