From b9245cca58f89cfb8d3a9dd30e0ae16b57c2099c Mon Sep 17 00:00:00 2001 From: jainapurva Date: Thu, 14 Nov 2024 14:58:53 -0800 Subject: [PATCH] Minor fixes --- test/hqq/test_hqq_affine.py | 2 +- torchao/dtypes/__init__.py | 2 -- torchao/prototype/hqq/example.py | 2 +- 3 files changed, 2 insertions(+), 4 deletions(-) diff --git a/test/hqq/test_hqq_affine.py b/test/hqq/test_hqq_affine.py index a71007518..2f231fbb3 100644 --- a/test/hqq/test_hqq_affine.py +++ b/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, ) diff --git a/torchao/dtypes/__init__.py b/torchao/dtypes/__init__.py index be1708be9..80b229bbb 100644 --- a/torchao/dtypes/__init__.py +++ b/torchao/dtypes/__init__.py @@ -1,6 +1,5 @@ from . import affine_quantized_tensor_ops -# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor from .affine_quantized_tensor import ( AffineQuantizedTensor, MarlinQQQTensor, @@ -29,7 +28,6 @@ PlainLayout, ) -# from ..prototype.dtypes.uint2 import UInt2Tensor, BitnetTensor __all__ = [ "NF4Tensor", diff --git a/torchao/prototype/hqq/example.py b/torchao/prototype/hqq/example.py index 0b562b05e..eb12b2b45 100644 --- a/torchao/prototype/hqq/example.py +++ b/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, )