fix rebase issue

Author: shiyang-weng

test/float8/test_compile.py

@@ -37,10 +37,6 @@
     hp_tensor_to_float8_dynamic,
 )
 from torchao.float8.float8_tensor import GemmInputRole, LinearMMConfig, ScaledMMConfig
-from torchao.quantization.quant_primitives import (
-    dequantize_affine_float8,
-    quantize_affine_float8,
-)
 from torchao.testing.float8.test_utils import get_test_float8_linear_config
 
 
@@ -412,14 +408,16 @@ def test_dynamic_scale_numeric_parity(
     ],
 )
 def test_quantize_dequantize_fp8_inductor(float8_dtype, hp_dtype):
+    quantize_affine_float8 = torch.ops.torchao.quantize_affine_float8
+    dequantize_affine_float8 = torch.ops.torchao.dequantize_affine_float8
     input = torch.randn(10, 10)
     with torch.no_grad():
         torch._dynamo.reset()
         expected_scale = torch.tensor(2.0)
         expected_quantized = quantize_affine_float8(
             input,
             expected_scale,
-            float8_dtype,
+            float8_dtype=float8_dtype,
         )
         expected_dequantized = dequantize_affine_float8(
             expected_quantized,
@@ -430,7 +428,7 @@ def test_quantize_dequantize_fp8_inductor(float8_dtype, hp_dtype):
             torch.compile(quantize_affine_float8),
             input,
             expected_scale,
-            float8_dtype,
+            float8_dtype=float8_dtype,
         )
         torch.testing.FileCheck().check(
             "torch.ops.torchao.quantize_affine_float8.default"

torchao/quantization/quant_primitives.py

@@ -2274,7 +2274,7 @@ def _expand_scale_to_tensor_shape(
 def _quantize_affine_float8(
     tensor: torch.Tensor,
     scale: torch.Tensor,
-    float8_dtype: torch.dtype,
+    float8_dtype: torch.dtype = torch.float8_e4m3fn,
 ) -> torch.Tensor:
     """
     Quantizes the high precision floating point tensor to a float8 tensor, using the given scaling factor.
@@ -2295,7 +2295,7 @@ def _quantize_affine_float8(
 def _quantize_affine_float8_meta(
     tensor: torch.Tensor,
     scale: torch.Tensor,
-    float8_dtype: torch.dtype,
+    float8_dtype: torch.dtype = torch.float8_e4m3fn,
 ) -> torch.Tensor:
     return torch.empty_like(tensor, dtype=float8_dtype)
 
@@ -2304,7 +2304,7 @@ def _quantize_affine_float8_meta(
 def _dequantize_affine_float8(
     tensor: torch.Tensor,
     scale: torch.Tensor,
-    output_dtype: torch.dtype,
+    output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
     """
     Dequantizes the float8 tensor to high precision tensor.
@@ -2322,6 +2322,6 @@ def _dequantize_affine_float8(
 def _dequantize_affine_float8_meta(
     tensor: torch.Tensor,
     scale: torch.Tensor,
-    output_dtype: torch.dtype,
+    output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
     return torch.empty_like(tensor, dtype=output_dtype)