refactor(qbits): remove subdirectory

Also avoid exporting AWQ and TinyGemm classes at the top level.
huggingface · Sep 20, 2024 · b011276 · b011276
1 parent 4f3af18
commit b011276
Show file tree

Hide file tree

Showing 13 changed files with 20 additions and 21 deletions.
diff --git a/...anto/tensor/weights/qbits/awq/__init__.py → ...mum/quanto/tensor/weights/awq/__init__.py b/...anto/tensor/weights/qbits/awq/__init__.py → ...mum/quanto/tensor/weights/awq/__init__.py
diff --git a/...quanto/tensor/weights/qbits/awq/packed.py → optimum/quanto/tensor/weights/awq/packed.py b/...quanto/tensor/weights/qbits/awq/packed.py → optimum/quanto/tensor/weights/awq/packed.py
diff --git a/.../quanto/tensor/weights/qbits/awq/qbits.py → optimum/quanto/tensor/weights/awq/qbits.py b/.../quanto/tensor/weights/qbits/awq/qbits.py → optimum/quanto/tensor/weights/awq/qbits.py
@@ -17,9 +17,9 @@
 import torch
 from torch.autograd import Function
 
-from ....function import QuantizedLinearFunction
-from ....grouped import group, ungroup
-from ....qtype import qtypes
+from ...function import QuantizedLinearFunction
+from ...grouped import group, ungroup
+from ...qtype import qtypes
 from ..qbits import WeightQBitsTensor
 from .packed import AWQPackedTensor, AWQPacking
 

diff --git a/optimum/quanto/tensor/weights/qbits/qbits.py → optimum/quanto/tensor/weights/qbits.py b/optimum/quanto/tensor/weights/qbits/qbits.py → optimum/quanto/tensor/weights/qbits.py
@@ -18,12 +18,12 @@
 from packaging import version
 from torch.autograd import Function
 
-from ...function import QuantizedLinearFunction
-from ...grouped import grouped_shape
-from ...packed import PackedTensor
-from ...qbits import QBitsTensor
-from ...qtensor import qfallback
-from ...qtype import qint2, qint4, qtype, qtypes
+from ..function import QuantizedLinearFunction
+from ..grouped import grouped_shape
+from ..packed import PackedTensor
+from ..qbits import QBitsTensor
+from ..qtensor import qfallback
+from ..qtype import qint2, qint4, qtype, qtypes
 
 
 __all__ = ["WeightQBitsTensor"]

diff --git a/optimum/quanto/tensor/weights/qbits/__init__.py b/optimum/quanto/tensor/weights/qbits/__init__.py
diff --git a/...tensor/weights/qbits/tinygemm/__init__.py → ...uanto/tensor/weights/tinygemm/__init__.py b/...tensor/weights/qbits/tinygemm/__init__.py → ...uanto/tensor/weights/tinygemm/__init__.py
diff --git a/...o/tensor/weights/qbits/tinygemm/packed.py → .../quanto/tensor/weights/tinygemm/packed.py b/...o/tensor/weights/qbits/tinygemm/packed.py → .../quanto/tensor/weights/tinygemm/packed.py
diff --git a/...to/tensor/weights/qbits/tinygemm/qbits.py → ...m/quanto/tensor/weights/tinygemm/qbits.py b/...to/tensor/weights/qbits/tinygemm/qbits.py → ...m/quanto/tensor/weights/tinygemm/qbits.py
@@ -17,10 +17,10 @@
 import torch
 from torch.autograd import Function
 
-from ....function import QuantizedLinearFunction
-from ....grouped import group, ungroup
-from ....qtype import qtypes
-from ...qbits import WeightQBitsTensor
+from ...function import QuantizedLinearFunction
+from ...grouped import group, ungroup
+from ...qtype import qtypes
+from ..qbits import WeightQBitsTensor
 from .packed import TinyGemmPackedTensor
 
 

diff --git a/test/library/test_mm.py b/test/library/test_mm.py
@@ -17,7 +17,7 @@
 import torch
 from helpers import assert_similar, random_tensor
 
-from optimum.quanto import AWQPackedTensor, AWQPacking
+from optimum.quanto.tensor.weights.awq import AWQPackedTensor, AWQPacking
 from optimum.quanto.tensor.weights.marlin.packed import get_scale_perms, pack_fp8_as_int32
 
 

diff --git a/test/tensor/weights/optimized/test_awq_packed_tensor.py b/test/tensor/weights/optimized/test_awq_packed_tensor.py
@@ -18,7 +18,7 @@
 import torch
 from helpers import device_eq
 
-from optimum.quanto import AWQPackedTensor, AWQPacking
+from optimum.quanto.tensor.weights.awq import AWQPackedTensor, AWQPacking
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

diff --git a/test/tensor/weights/optimized/test_awq_weight_qbits_tensor.py b/test/tensor/weights/optimized/test_awq_weight_qbits_tensor.py
@@ -17,7 +17,8 @@
 from helpers import device_eq, random_weight_qbits_tensor
 
 from optimum.quanto import qint4
-from optimum.quanto.tensor.weights import AWQWeightQBitsTensor, WeightQBitsTensor
+from optimum.quanto.tensor.weights.awq import AWQWeightQBitsTensor
+from optimum.quanto.tensor.weights import WeightQBitsTensor
 
 
 @pytest.mark.skipif(

diff --git a/test/tensor/weights/optimized/test_tinygemm_packed_tensor.py b/test/tensor/weights/optimized/test_tinygemm_packed_tensor.py
@@ -19,7 +19,7 @@
 from helpers import device_eq
 from packaging import version
 
-from optimum.quanto import TinyGemmPackedTensor
+from optimum.quanto.tensor.weights.tinygemm import TinyGemmPackedTensor
 
 
 @pytest.mark.skip_device("mps")  # Only available with pytorch 2.4

diff --git a/test/tensor/weights/optimized/test_tinygemm_weight_qbits_tensor.py b/test/tensor/weights/optimized/test_tinygemm_weight_qbits_tensor.py
@@ -18,7 +18,8 @@
 from packaging import version
 
 from optimum.quanto import qint4
-from optimum.quanto.tensor.weights import TinyGemmWeightQBitsTensor, WeightQBitsTensor
+from optimum.quanto.tensor.weights.tinygemm import TinyGemmWeightQBitsTensor
+from optimum.quanto.tensor.weights import WeightQBitsTensor
 
 
 @pytest.mark.skip_device("mps")  # Only available with pytorch 2.4