Minor fixes

Author: jainapurva

test/hqq/test_hqq_affine.py

@@ -1,6 +1,6 @@
 import unittest
 import torch
-from torchao.quantization.quant_primitives import (
+from torchao.quantization import (
     ZeroPointDomain,
     MappingType,
 )

torchao/dtypes/__init__.py

@@ -1,6 +1,4 @@
 from . import affine_quantized_tensor_ops
-
-# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
 from .affine_quantized_tensor import (
     AffineQuantizedTensor,
     MarlinQQQTensor,
@@ -29,8 +27,6 @@
     PlainLayout,
 )
 
-# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor
-
 __all__ = [
     "NF4Tensor",
     "to_nf4",

torchao/prototype/hqq/example.py

@@ -3,7 +3,7 @@
 from torchao.dtypes.affine_quantized_tensor import (
     to_affine_quantized_intx,
 )
-from torchao.quantization.quant_primitives import (
+from torchao.quantization import (
         ZeroPointDomain,
         MappingType,
 )