initial commit

README.md

@@ -1 +1 @@
-# sparsetensors
+# compressed-tensors

setup.py

@@ -18,7 +18,7 @@
 
 def _setup_packages() -> List:
     return find_packages(
-        "src", include=["sparsetensors", "sparsetensors.*"], exclude=["*.__pycache__.*"]
+        "src", include=["compressed-tensors", "compressed-tensors.*"], exclude=["*.__pycache__.*"]
     )
 
 def _setup_install_requires() -> List:
@@ -28,7 +28,7 @@ def _setup_extras() -> Dict:
     return {"dev": ["black==22.12.0", "isort==5.8.0", "wheel>=0.36.2", "flake8>=3.8.3", "pytest>=6.0.0",]}
 
 setup(
-    name="sparsetensors",
+    name="compressed-tensors",
     version="0.3.0",
     author="Neuralmagic, Inc.",
     author_email="support@neuralmagic.com",

src/sparsetensors/README.md

@@ -65,7 +65,7 @@ is saved to the model's `config.json` file. **Note:** the model must have been
 initialized with SparseAutoModelForCausalLM.from_pretrained()
 
 ```python
-from sparsetensors import BitmaskConfig
+from compressed_tensors import BitmaskConfig
 
 output_dir = "/PATH/TO/SAVE/COMPRESSED_MODEL"
 sparsity_config = BitmaskConfig()

src/sparsetensors/compressors/base.py

@@ -15,9 +15,9 @@
 import operator
 from typing import Dict, Generator, Tuple
 
-from sparsetensors.base import SPARSITY_CONFIG_NAME
-from sparsetensors.config import CompressionConfig
-from sparsetensors.registry import RegistryMixin
+from compressed_tensors.base import SPARSITY_CONFIG_NAME
+from compressed_tensors.config import CompressionConfig
+from compressed_tensors.registry import RegistryMixin
 from torch import Tensor
 from torch.nn import Module, Parameter
 from tqdm import tqdm

src/sparsetensors/compressors/dense.py

@@ -14,7 +14,7 @@
 
 from typing import Dict, Generator, Tuple
 
-from sparsetensors.compressors import ModelCompressor
+from compressed_tensors.compressors import ModelCompressor
 from torch import Tensor
 
 

src/sparsetensors/compressors/sparse_bitmask.py

@@ -18,8 +18,8 @@
 import numpy
 import torch
 from safetensors import safe_open
-from sparsetensors.compressors import ModelCompressor
-from sparsetensors.utils import get_nested_weight_mappings, merge_names
+from compressed_tensors.compressors import ModelCompressor
+from compressed_tensors.utils import get_nested_weight_mappings, merge_names
 from torch import Tensor
 from tqdm import tqdm
 

src/sparsetensors/config/base.py

@@ -15,7 +15,7 @@
 from typing import Optional
 
 from pydantic import BaseModel
-from sparsetensors.registry import RegistryMixin
+from compressed_tensors.registry import RegistryMixin
 
 
 __all__ = ["CompressionConfig"]

src/sparsetensors/config/dense.py

@@ -14,7 +14,7 @@
 
 from typing import Optional
 
-from sparsetensors.config import CompressionConfig
+from compressed_tensors.config import CompressionConfig
 
 
 __all__ = ["DenseSparsityConfig"]

src/sparsetensors/config/sparse_bitmask.py

@@ -14,7 +14,7 @@
 
 from typing import Optional
 
-from sparsetensors.config.base import CompressionConfig
+from compressed_tensors.config.base import CompressionConfig
 
 
 __all__ = ["BitmaskConfig"]

src/sparsetensors/quantization/lifecycle/apply.py

@@ -16,16 +16,16 @@
 from collections import OrderedDict
 from typing import Iterable, Optional
 
-from sparsetensors.quantization.lifecycle.calibration import set_module_for_calibration
-from sparsetensors.quantization.lifecycle.frozen import freeze_module_quantization
-from sparsetensors.quantization.lifecycle.initialize import (
+from compressed_tensors.quantization.lifecycle.calibration import set_module_for_calibration
+from compressed_tensors.quantization.lifecycle.frozen import freeze_module_quantization
+from compressed_tensors.quantization.lifecycle.initialize import (
     initialize_module_for_quantization,
 )
-from sparsetensors.quantization.quant_config import (
+from compressed_tensors.quantization.quant_config import (
     QuantizationConfig,
     QuantizationStatus,
 )
-from sparsetensors.quantization.utils import iter_named_leaf_modules
+from compressed_tensors.quantization.utils import iter_named_leaf_modules
 from torch.nn import Module
 
 

src/sparsetensors/quantization/lifecycle/calibration.py

@@ -15,7 +15,7 @@
 
 import logging
 
-from sparsetensors.quantization.quant_config import QuantizationStatus
+from compressed_tensors.quantization.quant_config import QuantizationStatus
 from torch.nn import Module
 
 

src/sparsetensors/quantization/lifecycle/forward.py

@@ -15,9 +15,9 @@
 from functools import wraps
 
 import torch
-from sparsetensors.quantization.quant_args import QuantizationArgs
-from sparsetensors.quantization.quant_config import QuantizationStatus
-from sparsetensors.quantization.quant_scheme import QuantizationScheme
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_config import QuantizationStatus
+from compressed_tensors.quantization.quant_scheme import QuantizationScheme
 from torch.nn import Module
 
 

src/sparsetensors/quantization/lifecycle/frozen.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from sparsetensors.quantization.quant_config import QuantizationStatus
+from compressed_tensors.quantization.quant_config import QuantizationStatus
 from torch.nn import Module
 
 

src/sparsetensors/quantization/lifecycle/initialize.py

@@ -17,10 +17,10 @@
 from typing import Optional
 
 import torch
-from sparsetensors.quantization.lifecycle.forward import wrap_module_forward_quantized
-from sparsetensors.quantization.quant_args import QuantizationArgs
-from sparsetensors.quantization.quant_config import QuantizationStatus
-from sparsetensors.quantization.quant_scheme import QuantizationScheme
+from compressed_tensors.quantization.lifecycle.forward import wrap_module_forward_quantized
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_config import QuantizationStatus
+from compressed_tensors.quantization.quant_scheme import QuantizationScheme
 from torch.nn import Module, Parameter
 
 

src/sparsetensors/quantization/observers/base.py

@@ -14,8 +14,8 @@
 
 from typing import Optional, Tuple
 
-from sparsetensors.quantization.quant_args import QuantizationArgs
-from sparsetensors.registry.registry import RegistryMixin
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.registry.registry import RegistryMixin
 from torch import FloatTensor, IntTensor, Tensor
 from torch.nn import Module
 

src/sparsetensors/quantization/observers/helpers.py

@@ -15,7 +15,7 @@
 from typing import Tuple
 
 import torch
-from sparsetensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_args import QuantizationArgs
 from torch import FloatTensor, IntTensor, Tensor
 
 

src/sparsetensors/quantization/observers/memoryless.py

@@ -15,8 +15,8 @@
 from typing import Tuple
 
 import torch
-from sparsetensors.quantization.observers.base import Observer
-from sparsetensors.quantization.observers.helpers import calculate_qparams
+from compressed_tensors.quantization.observers.base import Observer
+from compressed_tensors.quantization.observers.helpers import calculate_qparams
 from torch import FloatTensor, IntTensor, Tensor
 
 

src/sparsetensors/quantization/observers/min_max.py

@@ -15,9 +15,9 @@
 from typing import Tuple
 
 import torch
-from sparsetensors.quantization.observers.base import Observer
-from sparsetensors.quantization.observers.helpers import calculate_qparams
-from sparsetensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.observers.base import Observer
+from compressed_tensors.quantization.observers.helpers import calculate_qparams
+from compressed_tensors.quantization.quant_args import QuantizationArgs
 from torch import FloatTensor, IntTensor, Tensor
 
 

src/sparsetensors/quantization/quant_args.py

@@ -80,6 +80,6 @@ def get_observer(self):
         """
         :return: torch quantization FakeQuantize built based on these QuantizationArgs
         """
-        from sparsetensors.quantization.observers.base import Observer
+        from compressed_tensors.quantization.observers.base import Observer
 
         return Observer.load_from_registry(self.observer, quantization_args=self)

src/sparsetensors/quantization/quant_config.py

@@ -16,8 +16,8 @@
 from typing import Dict, List, Optional
 
 from pydantic import BaseModel, Field
-from sparsetensors.quantization.quant_scheme import QuantizationScheme
-from sparsetensors.quantization.utils import (
+from compressed_tensors.quantization.quant_scheme import QuantizationScheme
+from compressed_tensors.quantization.utils import (
     calculate_compression_ratio,
     is_module_quantized,
     iter_named_leaf_modules,

src/sparsetensors/quantization/quant_scheme.py

@@ -15,7 +15,7 @@
 from typing import List, Optional
 
 from pydantic import BaseModel
-from sparsetensors.quantization.quant_args import QuantizationArgs
+from compressed_tensors.quantization.quant_args import QuantizationArgs
 
 
 __all__ = ["QuantizationScheme"]

src/sparsetensors/utils/helpers.py

@@ -15,9 +15,9 @@
 
 from typing import Optional
 
-from sparsetensors.base import SPARSITY_CONFIG_NAME
-from sparsetensors.compressors import ModelCompressor
-from sparsetensors.config import CompressionConfig
+from compressed_tensors.base import SPARSITY_CONFIG_NAME
+from compressed_tensors.compressors import ModelCompressor
+from compressed_tensors.config import CompressionConfig
 from transformers import AutoConfig
 
 

tests/quantization/lifecycle/test_apply.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 
-from sparsetensors.quantization.lifecycle import apply_quantization_config
-from sparsetensors.quantization.quant_config import (
+from compressed_tensors.quantization.lifecycle import apply_quantization_config
+from compressed_tensors.quantization.quant_config import (
     QuantizationConfig,
     QuantizationStatus,
 )

tests/quantization/test_quant_args.py

@@ -14,7 +14,7 @@
 
 import pytest
 from pydantic import ValidationError
-from sparsetensors.quantization import (
+from compressed_tensors.quantization import (
     QuantizationArgs,
     QuantizationStrategy,
     QuantizationType,

tests/quantization/test_quant_config.py

@@ -15,7 +15,7 @@
 
 import pytest
 from pydantic import ValidationError
-from sparsetensors.quantization import (
+from compressed_tensors.quantization import (
     QuantizationConfig,
     QuantizationScheme,
     QuantizationStatus,

tests/quantization/test_quant_scheme.py

@@ -14,7 +14,7 @@
 
 import pytest
 from pydantic import ValidationError
-from sparsetensors.quantization import QuantizationArgs, QuantizationScheme
+from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme
 
 
 def test_basic_scheme():

tests/test_bitmask.py

@@ -18,7 +18,7 @@
 import pytest
 import torch
 from safetensors.torch import save_file
-from sparsetensors import BitmaskCompressor, BitmaskConfig, BitmaskTensor
+from compressed_tensors import BitmaskCompressor, BitmaskConfig, BitmaskTensor
 
 
 @pytest.mark.parametrize(

tests/test_registry.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import pytest
-from sparsetensors import (
+from compressed_tensors import (
     BitmaskCompressor,
     BitmaskConfig,
     CompressionConfig,