Merge branch 'main' into wengshiy/int8_scaled_embedding_bag

Author: shiyang-weng

test/quantization/pt2e/test_quantize_pt2e.py

@@ -2948,11 +2948,10 @@ def has_inplace_ops(graph_module: torch.fx.GraphModule) -> bool:
 @unittest.skipIf(not torch_version_at_least("2.7.0"), "Requires torch 2.7+")
 class TestQuantizePT2EAffineQuantization(PT2EQuantizationTestCase):
     def test_channel_group_quantization(self):
-        from torchao.quantization import PerGroup, PerToken
         from torchao.quantization.pt2e._affine_quantization import (
             AffineQuantizedMinMaxObserver,
         )
-        from torchao.quantization.pt2e.observer import MappingType
+        from torchao.quantization.pt2e.observer import MappingType, PerGroup, PerToken
 
         class BackendAQuantizer(Quantizer):
             def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
@@ -3032,13 +3031,13 @@ def forward(self, x):
     def test_dynamic_affine_act_per_channel_weights(self):
         import operator
 
-        from torchao.quantization import PerToken
         from torchao.quantization.pt2e._affine_quantization import (
             AffineQuantizedMovingAverageMinMaxObserver,
         )
         from torchao.quantization.pt2e.observer import (
             MappingType,
             PerChannelMinMaxObserver,
+            PerToken,
         )
 
         class BackendAQuantizer(Quantizer):
@@ -3123,14 +3122,12 @@ def forward(self, x):
     def test_dynamic_per_tok_act_per_group_weights(self):
         import operator
 
-        from torchao.quantization import PerGroup, PerToken
-
         # TODO: merge into torchao observer
         from torchao.quantization.pt2e._affine_quantization import (
             AffineQuantizedMinMaxObserver,
             AffineQuantizedPlaceholderObserver,
         )
-        from torchao.quantization.pt2e.observer import MappingType
+        from torchao.quantization.pt2e.observer import MappingType, PerGroup, PerToken
 
         class BackendAQuantizer(Quantizer):
             def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:

test/test_ops.py

@@ -40,7 +40,14 @@
 except RuntimeError:
     pytest.skip("torchao.ops not available")
 
+from torchao.quantization import PerGroup, PerRow, PerTensor
+from torchao.quantization.quant_primitives import (
+    _choose_scale_float8,
+    _dequantize_affine_float8,
+    _quantize_affine_float8,
+)
 from torchao.quantization.utils import (
+    get_block_size,
     get_groupwise_affine_qparams,
     groupwise_affine_dequantize_tensor_from_qparams,
     groupwise_affine_quantize_tensor_from_qparams,