Suppress deprecation warnings from #4074 and #4100 (#4155)

**Summary:** #4074 and #4100 deprecated a few classes, but this
triggered the following warnings when the user imports torchao
from the top-level. This commit suppresses these warnings
in this case.

Before:
```
import torchao
/data/users/andrewor/ao/torchao/dtypes/utils.py:89: UserWarning: Deprecation: PlainLayout is deprecated and will be removed in a future release of torchao, see #2752 for more details
  warnings.warn(
/data/users/andrewor/ao/torchao/quantization/quant_primitives.py:95: UserWarning: Deprecation: TorchAODType is deprecated, please use the torch.intN dtype instead (e.g. TorchAODType.INT4 -> torch.int4)
  warnings.warn(
/data/users/andrewor/ao/torchao/dtypes/utils.py:89: UserWarning: Deprecation: PlainLayout is deprecated and will be removed in a future release of torchao, see #2752 for more details
  warnings.warn(
```

After:
```
import torchao
\# No warnings
```

**Test Plan:** Manual testing.

Author: andrewor14

torchao/dtypes/affine_quantized_tensor.py

@@ -234,7 +234,7 @@ def from_hp_to_intx(
         zero_point_dtype: Optional[torch.dtype] = None,
         preserve_zero: bool = True,
         zero_point_domain: ZeroPointDomain = _DEFAULT_ZPD,  # type: ignore[assignment]
-        _layout: Layout = PlainLayout(),
+        _layout: Optional[Layout] = None,
         use_hqq: bool = False,
         *,
         custom_scale: Optional[torch.Tensor] = None,
@@ -253,6 +253,8 @@ def from_hp_to_intx(
             quantize_affine,
         )
 
+        if _layout is None:
+            _layout = PlainLayout()
         if zero_point_domain is _DEFAULT_ZPD:
             zero_point_domain = ZeroPointDomain.INT
 
@@ -399,7 +401,7 @@ def from_hp_to_intx_static(
         quant_min: Optional[int] = None,
         quant_max: Optional[int] = None,
         zero_point_domain: ZeroPointDomain = _DEFAULT_ZPD,  # type: ignore[assignment]
-        _layout: Layout = PlainLayout(),
+        _layout: Optional[Layout] = None,
     ):
         """Create an integer AffineQuantizedTensor from a high precision tensor using static parameters."""
         from torchao.quantization.quant_primitives import (
@@ -409,6 +411,8 @@ def from_hp_to_intx_static(
             quantize_affine,
         )
 
+        if _layout is None:
+            _layout = PlainLayout()
         if zero_point_domain is _DEFAULT_ZPD:
             zero_point_domain = ZeroPointDomain.INT
         if zero_point_domain is None:

torchao/quantization/qat/fake_quantize_config.py

@@ -109,8 +109,6 @@ class IntxFakeQuantizeConfig(FakeQuantizeConfigBase):
 
     Args:
         dtype: dtype to simulate during fake quantization, e.g. torch.int8.
-            For PyTorch versions older than 2.6, you may use `TorchAODType` to represent
-            torch.int1 to torch.int7 instead, e.g. TorchAODType.INT4.
         granularity: granularity of scales and zero points, e.g. PerGroup(32).
             We also support the following strings:
                1) 'per_token': equivalent to PerToken()
@@ -151,7 +149,7 @@ class IntxFakeQuantizeConfig(FakeQuantizeConfigBase):
         IntxFakeQuantizeConfig(torch.int4, PerGroup(32), MappingType.SYMMETRIC)
     """
 
-    dtype: Union[torch.dtype, TorchAODType]
+    dtype: Union[torch.dtype, "TorchAODType"]
     granularity: Granularity
     mapping_type: MappingType
     scale_precision: torch.dtype
@@ -163,7 +161,7 @@ class IntxFakeQuantizeConfig(FakeQuantizeConfigBase):
 
     def __init__(
         self,
-        dtype: Union[torch.dtype, TorchAODType],
+        dtype: Union[torch.dtype, "TorchAODType"],
         granularity: Union[Granularity, str, None] = None,
         mapping_type: Optional[MappingType] = None,
         scale_precision: torch.dtype = torch.float32,

torchao/quantization/quant_api.py

@@ -1095,7 +1095,7 @@ class Int8DynamicActivationInt8WeightConfig(AOBaseConfig):
        :language: python
     """
 
-    layout: Optional[Layout] = PlainLayout()
+    layout: Optional[Layout] = None
     act_mapping_type: Optional[MappingType] = MappingType.SYMMETRIC
     weight_only_decode: bool = False
     granularity: Optional[
@@ -1108,6 +1108,8 @@ def __post_init__(self):
         torch._C._log_api_usage_once(
             "torchao.quantization.Int8DynamicActivationInt8WeightConfig"
         )
+        if self.layout is None:
+            self.layout = PlainLayout()
         if self.version == 2:
             act_granularity, weight_granularity = Int8Tensor._normalize_granularity(
                 self.granularity

torchao/quantization/quant_primitives.py

@@ -92,10 +92,11 @@ class ZeroPointDomain(Enum):
 class _TorchAODTypeMeta(EnumMeta):
     def __getattribute__(cls, name):
         result = super().__getattribute__(name)
-        warnings.warn(
-            "Deprecation: TorchAODType is deprecated, please use the torch.intN dtype instead "
-            "(e.g. TorchAODType.INT4 -> torch.int4)"
-        )
+        if isinstance(result, cls):
+            warnings.warn(
+                "Deprecation: TorchAODType is deprecated, please use the torch.intN dtype instead "
+                "(e.g. TorchAODType.INT4 -> torch.int4)"
+            )
         return result
 
 
@@ -128,36 +129,41 @@ class TorchAODType(Enum, metaclass=_TorchAODTypeMeta):
 Map from dtype to the bound value of integers
 TODO: maybe can replace this with call to torch.iinfo
 """
-_DTYPE_TO_QVALUE_BOUNDS: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = {
-    torch.uint8: (0, 255),
-    torch.int8: (-128, 127),
-    torch.int16: (-(2**15), 2**15 - 1),
-    torch.int32: (-(2**31), 2**31 - 1),
-}
-_DTYPE_TO_BIT_WIDTH: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = {
-    TorchAODType.INT1: 1,
-    TorchAODType.INT2: 2,
-    TorchAODType.INT3: 3,
-    TorchAODType.INT4: 4,
-    TorchAODType.INT5: 5,
-    TorchAODType.INT6: 6,
-    TorchAODType.INT7: 7,
-    torch.uint8: 8,
-    torch.int8: 8,
-    torch.int16: 16,
-    torch.int32: 32,
-}
+# Suppress TorchAODType deprecation warnings for internal usage
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore", UserWarning)
+
+    _DTYPE_TO_QVALUE_BOUNDS: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = {
+        torch.uint8: (0, 255),
+        torch.int8: (-128, 127),
+        torch.int16: (-(2**15), 2**15 - 1),
+        torch.int32: (-(2**31), 2**31 - 1),
+    }
 
-_SUB_BYTE_UINT_BOUNDS: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = {}
-_SUB_BYTE_INT_BOUNDS: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = {
-    TorchAODType.INT1: (-(2**0), 2**0 - 1),
-    TorchAODType.INT2: (-(2**1), 2**1 - 1),
-    TorchAODType.INT3: (-(2**2), 2**2 - 1),
-    TorchAODType.INT4: (-(2**3), 2**3 - 1),
-    TorchAODType.INT5: (-(2**4), 2**4 - 1),
-    TorchAODType.INT6: (-(2**5), 2**5 - 1),
-    TorchAODType.INT7: (-(2**6), 2**6 - 1),
-}
+    _DTYPE_TO_BIT_WIDTH: Dict[Union[torch.dtype, TorchAODType], int] = {
+        TorchAODType.INT1: 1,
+        TorchAODType.INT2: 2,
+        TorchAODType.INT3: 3,
+        TorchAODType.INT4: 4,
+        TorchAODType.INT5: 5,
+        TorchAODType.INT6: 6,
+        TorchAODType.INT7: 7,
+        torch.uint8: 8,
+        torch.int8: 8,
+        torch.int16: 16,
+        torch.int32: 32,
+    }
+
+    _SUB_BYTE_UINT_BOUNDS: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = {}
+    _SUB_BYTE_INT_BOUNDS: Dict[Union[torch.dtype, TorchAODType], Tuple[int, int]] = {
+        TorchAODType.INT1: (-(2**0), 2**0 - 1),
+        TorchAODType.INT2: (-(2**1), 2**1 - 1),
+        TorchAODType.INT3: (-(2**2), 2**2 - 1),
+        TorchAODType.INT4: (-(2**3), 2**3 - 1),
+        TorchAODType.INT5: (-(2**4), 2**4 - 1),
+        TorchAODType.INT6: (-(2**5), 2**5 - 1),
+        TorchAODType.INT7: (-(2**6), 2**6 - 1),
+    }
 
 _SUB_BYTE_UINT_BOUNDS = {
     torch.uint1: (0, 2**1 - 1),