compressed-tensors 0.13.1a20260130__tar.gz → 0.13.1a20260203__tar.gz

{compressed_tensors-0.13.1a20260130/src/compressed_tensors.egg-info → compressed_tensors-0.13.1a20260203}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: compressed-tensors
-Version: 0.13.1a20260130
+Version: 0.13.1a20260203
 Summary: Library for utilization of compressed safetensors of neural network models
 Home-page: https://github.com/vllm-project/compressed-tensors
 Author: Neuralmagic, Inc.

{compressed_tensors-0.13.1a20260130 → compressed_tensors-0.13.1a20260203}/src/compressed_tensors/compressors/base.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Dict, Generator, Optional, Tuple, Union
+from collections.abc import Generator
 
 import torch
 from compressed_tensors.config import SparsityCompressionConfig
@@ -59,15 +59,15 @@ class BaseCompressor(RegistryMixin, ABC):
     """
 
     def __init__(
-        self, config: Union[SparsityCompressionConfig, QuantizationConfig, None] = None
+        self, config: SparsityCompressionConfig | QuantizationConfig | None = None
     ):
         self.config = config
 
     def compression_param_info(
         self,
         weight_shape: torch.Size,
-        quantization_args: Optional[QuantizationArgs] = None,
-    ) -> Dict[str, Tuple[torch.Size, torch.dtype]]:
+        quantization_args: QuantizationArgs | None = None,
+    ) -> dict[str, tuple[torch.Size, torch.dtype]]:
         """
         Creates a dictionary of expected shapes and dtypes for each compression
             parameter used by the compressor
@@ -80,7 +80,7 @@ class BaseCompressor(RegistryMixin, ABC):
 
     @property
     @abstractmethod
-    def compression_param_names(self) -> Tuple[str]:
+    def compression_param_names(self) -> tuple[str, ...]:
         """
         Returns a tuple of compression parameter names introduced by
         the compressor during compression
@@ -90,9 +90,9 @@ class BaseCompressor(RegistryMixin, ABC):
     @abstractmethod
     def compress(
         self,
-        model_state: Dict[str, Tensor],
+        model_state: dict[str, Tensor],
         **kwargs,
-    ) -> Dict[str, Tensor]:
+    ) -> dict[str, Tensor]:
         """
         Compresses a dense state dict
 
@@ -108,7 +108,7 @@ class BaseCompressor(RegistryMixin, ABC):
         path_to_model_or_tensors: str,
         device: str = "cpu",
         **kwargs,
-    ) -> Generator[Tuple[str, Tensor], None, None]:
+    ) -> Generator[tuple[str, Tensor], None, None]:
         """
         Reads a compressed state dict located at path_to_model_or_tensors
         and returns a generator for sequentially decompressing back to a
@@ -122,7 +122,7 @@ class BaseCompressor(RegistryMixin, ABC):
         """
         raise NotImplementedError()
 
-    def compress_module(self, module: Module) -> Optional[Dict[str, torch.Tensor]]:
+    def compress_module(self, module: Module) -> dict[str, torch.Tensor] | None:
         """
         Compresses a single quantized leaf PyTorch module. If the module is not
         quantized, this function has no effect.
@@ -153,7 +153,7 @@ class BaseCompressor(RegistryMixin, ABC):
         self,
         weight: Tensor,
         **kwargs,
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         """
         Compresses a single uncompressed weight
 
@@ -196,7 +196,7 @@ class BaseCompressor(RegistryMixin, ABC):
         return decompressed_weight
 
     def decompress_weight(
-        self, compressed_data: Dict[str, Tensor], **kwargs
+        self, compressed_data: dict[str, Tensor], **kwargs
     ) -> torch.Tensor:
         """
         Decompresses a single compressed weight

{compressed_tensors-0.13.1a20260130 → compressed_tensors-0.13.1a20260203}/src/compressed_tensors/compressors/helpers.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections.abc import Generator
 from pathlib import Path
-from typing import Dict, Generator, Optional, Tuple, Union
 
 import torch
 from compressed_tensors.compressors import BaseCompressor
@@ -32,9 +32,9 @@ __all__ = [
 
 
 def save_compressed(
-    tensors: Dict[str, Tensor],
-    save_path: Union[str, Path],
-    compression_format: Optional[CompressionFormat] = None,
+    tensors: dict[str, Tensor],
+    save_path: str | Path,
+    compression_format: CompressionFormat | None = None,
 ):
     """
     Save compressed tensors to disk. If tensors are not compressed,
@@ -68,10 +68,10 @@ def save_compressed(
 
 
 def load_compressed(
-    compressed_tensors: Union[str, Path],
+    compressed_tensors: str | Path,
     compression_config: SparsityCompressionConfig = None,
-    device: Optional[str] = "cpu",
-) -> Generator[Tuple[str, Tensor], None, None]:
+    device: str | None = "cpu",
+) -> Generator[tuple[str, Tensor], None, None]:
     """
     Load compressed tensors from disk.
     If tensors are not compressed, load them as is.
@@ -111,7 +111,7 @@ def load_compressed(
 def save_compressed_model(
     model: torch.nn.Module,
     filename: str,
-    compression_format: Optional[CompressionFormat] = None,
+    compression_format: CompressionFormat | None = None,
     force_contiguous: bool = True,
 ):
     """

{compressed_tensors-0.13.1a20260130 → compressed_tensors-0.13.1a20260203}/src/compressed_tensors/compressors/model_compressors/model_compressor.py

@@ -18,7 +18,7 @@ import operator
 import os
 import re
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, TypeVar, Union
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import compressed_tensors
 import torch
@@ -109,16 +109,16 @@ class ModelCompressor:
     :param quantization_config: config specifying quantization compression parameters
     """
 
-    sparsity_config: Optional[SparsityCompressionConfig] = None
-    quantization_config: Optional[QuantizationConfig] = None
-    transform_config: Optional[TransformConfig] = None
+    sparsity_config: SparsityCompressionConfig | None = None
+    quantization_config: QuantizationConfig | None = None
+    transform_config: TransformConfig | None = None
 
     @classmethod
     def from_pretrained(
         cls,
         pretrained_model_name_or_path: str,
         **kwargs,
-    ) -> Optional["ModelCompressor"]:
+    ) -> "ModelCompressor | None":
         """
         Given a path to a model config, extract the sparsity and/or quantization
         configs and load a ModelCompressor
@@ -133,7 +133,7 @@ class ModelCompressor:
     @classmethod
     def from_compression_config(
         cls,
-        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"],
+        compression_config: "dict[str, Any] | CompressedTensorsConfig",
     ):
         """
         :param compression_config:
@@ -172,10 +172,10 @@ class ModelCompressor:
     def from_pretrained_model(
         cls,
         model: Module,
-        sparsity_config_or_format: Union[SparsityCompressionConfig, str, None] = None,
-        quantization_format: Optional[str] = None,
-        sparsity_config: Union[SparsityCompressionConfig, str, None] = None,
-    ) -> Optional["ModelCompressor"]:
+        sparsity_config_or_format: SparsityCompressionConfig | str | None = None,
+        quantization_format: str | None = None,
+        sparsity_config: SparsityCompressionConfig | str | None = None,
+    ) -> "ModelCompressor | None":
         """
         Given a pytorch model and optional sparsity and/or quantization configs,
         load the appropriate compressors
@@ -232,8 +232,8 @@ class ModelCompressor:
 
     @staticmethod
     def parse_sparsity_config(
-        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"],
-    ) -> Union[Dict[str, Any], None]:
+        compression_config: "dict[str, Any] | CompressedTensorsConfig",
+    ) -> dict[str, Any] | None:
         """
         Parse sparsity config from quantization/compression config. Sparsity
         config is nested inside q/c config
@@ -253,8 +253,8 @@ class ModelCompressor:
 
     @staticmethod
     def parse_quantization_config(
-        compression_config: Union[Dict[str, Any], "CompressedTensorsConfig"],
-    ) -> Union[Dict[str, Any], None]:
+        compression_config: "dict[str, Any] | CompressedTensorsConfig",
+    ) -> dict[str, Any] | None:
         """
         Parse quantization config from quantization/compression config. The
         quantization are all the fields that are not the sparsity config or
@@ -289,7 +289,7 @@ class ModelCompressor:
 
         return quantization_config
 
-    def _fetch_unique_quantization_formats(self) -> List[str]:
+    def _fetch_unique_quantization_formats(self) -> list[str]:
         """
         Get all unique compression formats present in a model.
         :return: list of quantization formats
@@ -309,10 +309,10 @@ class ModelCompressor:
 
     def __init__(
         self,
-        sparsity_config: Optional[SparsityCompressionConfig] = None,
-        quantization_config: Optional[QuantizationConfig] = None,
-        transform_config: Optional[TransformConfig] = None,
-        compression_formats: Optional[List[str]] = None,
+        sparsity_config: SparsityCompressionConfig | None = None,
+        quantization_config: QuantizationConfig | None = None,
+        transform_config: TransformConfig | None = None,
+        compression_formats: list[str] | None = None,
     ):
         self.sparsity_config = sparsity_config
         self.quantization_config = quantization_config
@@ -320,8 +320,8 @@ class ModelCompressor:
         self.compression_formats = compression_formats
 
         self.sparsity_compressor = None
-        self.quantization_compressor: Optional[
-            Dict[str, Union[BaseQuantizationCompressor, DenseCompressor]]
+        self.quantization_compressor: dict[
+            str, BaseQuantizationCompressor | DenseCompressor
         ] = None
         # no transform compressor is required
 
@@ -345,7 +345,7 @@ class ModelCompressor:
                     format, config=quantization_config
                 )
 
-    def get_missing_module_keys(self, model: Module) -> List[str]:
+    def get_missing_module_keys(self, model: Module) -> list[str]:
         """
         Identifies the expected missing weight keys in the compressed state_dict.
 
@@ -394,7 +394,7 @@ class ModelCompressor:
 
         return list(missing_keys)
 
-    def get_unexpected_file_keys(self, model: Module) -> List[str]:
+    def get_unexpected_file_keys(self, model: Module) -> list[str]:
         """
         Identifies extra keys introduced by the compression process in the
         compressed state_dict that are not expected by the model graph.
@@ -625,9 +625,9 @@ class ModelCompressor:
     def compress(
         self,
         model: Module,
-        state_dict: Optional[Dict[str, Tensor]] = None,
+        state_dict: dict[str, Tensor] | None = None,
         show_progress: bool = False,
-    ) -> Dict[str, Tensor]:
+    ) -> dict[str, Tensor]:
         """
         Compresses a dense state dict or model with sparsity and/or quantization
 
@@ -656,7 +656,7 @@ class ModelCompressor:
                 )
 
         if self.sparsity_compressor is not None:
-            sparse_compression_targets: Set[str] = {
+            sparse_compression_targets: set[str] = {
                 module_name
                 for module_name, _module in match_named_modules(
                     model=model,
@@ -732,7 +732,7 @@ class ModelCompressor:
                 QuantizationStatus.FROZEN,
             ):
                 apply_quantization_config(model, self.quantization_config)
-                names_to_scheme: Set[QuantizationScheme] = {
+                names_to_scheme: dict[str, QuantizationScheme] = {
                     name: getattr(module, "quantization_scheme")
                     for name, module in model.named_modules()
                     if getattr(module, "quantization_scheme", None) is not None
@@ -897,7 +897,7 @@ class ModelCompressor:
                         update_parameter_data(module, param_data, param_name)
 
 
-def map_module_to_scheme(model: Module) -> Dict[str, QuantizationScheme]:
+def map_module_to_scheme(model: Module) -> dict[str, QuantizationScheme]:
     """
     Returns a dictionary which maps quantized module names to their quantization
     schemes. Only includes modules with weight quantization

{compressed_tensors-0.13.1a20260130 → compressed_tensors-0.13.1a20260203}/src/compressed_tensors/compressors/quantized_compressors/base.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 import logging
+from collections.abc import Generator
 from pathlib import Path
-from typing import Any, Dict, Generator, Tuple, Union
+from typing import Any
 
 import torch
 from compressed_tensors.compressors.base import BaseCompressor
@@ -68,12 +69,12 @@ class BaseQuantizationCompressor(BaseCompressor):
 
     def compress(
         self,
-        model_state: Dict[str, Tensor],
-        names_to_scheme: Dict[str, QuantizationScheme],
+        model_state: dict[str, Tensor],
+        names_to_scheme: dict[str, QuantizationScheme],
         show_progress: bool = False,
         compression_device: str = "cpu",
         **kwargs,
-    ) -> Dict[str, Tensor]:
+    ) -> dict[str, Tensor]:
         """
         Compresses a dense state dict
 
@@ -141,7 +142,7 @@ class BaseQuantizationCompressor(BaseCompressor):
         return isinstance(self, NVFP4PackedCompressor)
 
     def _skip_zp(
-        self, name: str, names_to_scheme: Dict[str, QuantizationScheme]
+        self, name: str, names_to_scheme: dict[str, QuantizationScheme]
     ) -> bool:
         from compressed_tensors.compressors import PackedQuantizationCompressor
 
@@ -169,10 +170,10 @@ class BaseQuantizationCompressor(BaseCompressor):
 
     def decompress(
         self,
-        path_to_model_or_tensors: Union[str, Path, Dict[str, Any]],
-        names_to_scheme: Dict[str, QuantizationScheme],
+        path_to_model_or_tensors: str | Path | dict[str, Any],
+        names_to_scheme: dict[str, QuantizationScheme],
         device: str = "cpu",
-    ) -> Generator[Tuple[str, Tensor], None, None]:
+    ) -> Generator[tuple[str, Tensor], None, None]:
         """
         Reads a compressed state dict located at path_to_model_or_tensors
         and returns a generator for sequentially decompressing back to a
@@ -196,8 +197,8 @@ class BaseQuantizationCompressor(BaseCompressor):
 
     def _decompress_from_path(
         self,
-        path_to_model: Union[str, Path, Dict[str, Any]],
-        names_to_scheme: Dict[str, QuantizationScheme],
+        path_to_model: str | Path | dict[str, Any],
+        names_to_scheme: dict[str, QuantizationScheme],
         device: str,
     ):
         weight_mappings = get_nested_weight_mappings(
@@ -219,9 +220,9 @@ class BaseQuantizationCompressor(BaseCompressor):
 
     def decompress_from_state_dict(
         self,
-        state_dict: Dict[str, torch.Tensor],
-        names_to_scheme: Dict[str, QuantizationScheme],
-    ) -> Generator[Tuple[str, Dict[str, torch.Tensor]], None, None]:
+        state_dict: dict[str, torch.Tensor],
+        names_to_scheme: dict[str, QuantizationScheme],
+    ) -> Generator[tuple[str, dict[str, torch.Tensor]], None, None]:
         weight_mappings = get_nested_mappings_from_state_dict(
             state_dict, self.compression_param_names
         )
@@ -239,9 +240,9 @@ class BaseQuantizationCompressor(BaseCompressor):
     def decompress_module_from_state_dict(
         self,
         prefix: str,
-        state_dict: Dict[str, torch.Tensor],
+        state_dict: dict[str, torch.Tensor],
         scheme: QuantizationScheme,
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         """
         Only used by in-memory decompression pathways to decompress the parameters of
         one module

{compressed_tensors-0.13.1a20260130 → compressed_tensors-0.13.1a20260203}/src/compressed_tensors/compressors/quantized_compressors/fp4_quantized.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 
 
-from typing import Dict, Optional, Tuple
-
 import torch
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.compressors.quantized_compressors.base import (
@@ -48,7 +46,7 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
     """
 
     @property
-    def compression_param_names(self) -> Tuple[str]:
+    def compression_param_names(self) -> tuple[str, ...]:
         """
         Returns a tuple of compression parameter names introduced by
         the compressor during compression
@@ -63,8 +61,8 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
     def compression_param_info(
         self,
         weight_shape: torch.Size,
-        quantization_args: Optional[QuantizationArgs] = None,
-    ) -> Dict[str, Tuple[torch.Size, torch.dtype]]:
+        quantization_args: QuantizationArgs | None = None,
+    ) -> dict[str, tuple[torch.Size, torch.dtype]]:
         """
         Creates a dictionary of expected shapes and dtypes for each compression
             parameter used by the compressor
@@ -85,7 +83,7 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
         self,
         scale: Tensor,
         quantization_args: QuantizationArgs,
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         assert quantization_args.scale_dtype is not None
         return scale.to(quantization_args.scale_dtype)
 
@@ -95,10 +93,10 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
         scale: Tensor,
         global_scale: Tensor,
         quantization_args: QuantizationArgs,
-        device: Optional[torch.device] = None,
-        zero_point: Optional[torch.Tensor] = None,
-        g_idx: Optional[torch.Tensor] = None,
-    ) -> Dict[str, torch.Tensor]:
+        device: torch.device | None = None,
+        zero_point: torch.Tensor | None = None,
+        g_idx: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
         quantized_weight = quantize(
             x=weight,
             scale=scale,
@@ -118,8 +116,8 @@ class NVFP4PackedCompressor(BaseQuantizationCompressor):
 
     def decompress_weight(
         self,
-        compressed_data: Dict[str, Tensor],
-        quantization_args: Optional[QuantizationArgs] = None,
+        compressed_data: dict[str, Tensor],
+        quantization_args: QuantizationArgs | None = None,
     ) -> torch.Tensor:
         weight = compressed_data["weight_packed"]
         global_scale = compressed_data["weight_global_scale"]
@@ -149,15 +147,15 @@ class MXFP4PackedCompressor(NVFP4PackedCompressor):
         self,
         scale: Tensor,
         quantization_args: QuantizationArgs,
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         assert quantization_args.scale_dtype is not None
         scale_exp = 127 + torch.floor(torch.log2(scale)).to(torch.int32)
         return scale_exp.to(quantization_args.scale_dtype)
 
     def decompress_weight(
         self,
-        compressed_data: Dict[str, Tensor],
-        quantization_args: Optional[QuantizationArgs] = None,
+        compressed_data: dict[str, Tensor],
+        quantization_args: QuantizationArgs | None = None,
     ) -> torch.Tensor:
         raise NotImplementedError("MXFP4 Decompression is currently not supported")
 
@@ -216,7 +214,7 @@ kE2M1ToFloat = torch.tensor(
 # reference: : https://github.com/vllm-project/vllm/pull/16362
 @torch.compile(fullgraph=True, dynamic=True)
 def unpack_fp4_from_uint8(
-    a: torch.Tensor, m: int, n: int, dtype: Optional[torch.dtype] = torch.bfloat16
+    a: torch.Tensor, m: int, n: int, dtype: torch.dtype | None = torch.bfloat16
 ) -> torch.Tensor:
     """
     Unpacks uint8 values into fp4. Each uint8 consists of two fp4 values

{compressed_tensors-0.13.1a20260130 → compressed_tensors-0.13.1a20260203}/src/compressed_tensors/compressors/quantized_compressors/naive_quantized.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Optional, Tuple
-
 import torch
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.compressors.quantized_compressors.base import (
@@ -42,7 +40,7 @@ class NaiveQuantizationCompressor(BaseQuantizationCompressor):
     """
 
     @property
-    def compression_param_names(self) -> Tuple[str]:
+    def compression_param_names(self) -> tuple[str, ...]:
         """
         Returns a tuple of compression parameter names introduced by
         the compressor during compression
@@ -57,8 +55,8 @@ class NaiveQuantizationCompressor(BaseQuantizationCompressor):
     def compression_param_info(
         self,
         weight_shape: torch.Size,
-        quantization_args: Optional[QuantizationArgs] = None,
-    ) -> Dict[str, Tuple[torch.Size, torch.dtype]]:
+        quantization_args: QuantizationArgs | None = None,
+    ) -> dict[str, tuple[torch.Size, torch.dtype]]:
         """
         Creates a dictionary of expected shapes and dtypes for each compression
             parameter used by the compressor
@@ -75,11 +73,11 @@ class NaiveQuantizationCompressor(BaseQuantizationCompressor):
         weight: Tensor,
         scale: Tensor,
         quantization_args: QuantizationArgs,
-        zero_point: Optional[Tensor] = None,
-        g_idx: Optional[torch.Tensor] = None,
-        device: Optional[torch.device] = None,
-        global_scale: Optional[torch.Tensor] = None,
-    ) -> Dict[str, torch.Tensor]:
+        zero_point: Tensor | None = None,
+        g_idx: torch.Tensor | None = None,
+        device: torch.device | None = None,
+        global_scale: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
         """
         Compresses a single uncompressed weight
 
@@ -115,8 +113,8 @@ class NaiveQuantizationCompressor(BaseQuantizationCompressor):
 
     def decompress_weight(
         self,
-        compressed_data: Dict[str, Tensor],
-        quantization_args: Optional[QuantizationArgs] = None,
+        compressed_data: dict[str, Tensor],
+        quantization_args: QuantizationArgs | None = None,
     ) -> torch.Tensor:
         """
         Decompresses a single compressed weight

{compressed_tensors-0.13.1a20260130 → compressed_tensors-0.13.1a20260203}/src/compressed_tensors/compressors/quantized_compressors/pack_quantized.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-from typing import Dict, Literal, Optional, Tuple, Union
+from typing import Literal
 
 import torch
 from compressed_tensors.compressors.base import BaseCompressor
@@ -36,7 +36,7 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
     """
 
     @property
-    def compression_param_names(self) -> Tuple[str]:
+    def compression_param_names(self) -> tuple[str, ...]:
         """
         Returns a tuple of compression parameter names introduced by
         the compressor during compression
@@ -52,8 +52,8 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
     def compression_param_info(
         self,
         weight_shape: torch.Size,
-        quantization_args: Optional[QuantizationArgs] = None,
-    ) -> Dict[str, Tuple[torch.Size, torch.dtype]]:
+        quantization_args: QuantizationArgs | None = None,
+    ) -> dict[str, tuple[torch.Size, torch.dtype]]:
         """
         Creates a dictionary of expected shapes and dtypes for each compression
             parameter used by the compressor
@@ -90,11 +90,11 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
         weight: Tensor,
         scale: Tensor,
         quantization_args: QuantizationArgs,
-        zero_point: Optional[Tensor] = None,
-        g_idx: Optional[torch.Tensor] = None,
-        device: Optional[torch.device] = None,
-        global_scale: Optional[torch.Tensor] = None,
-    ) -> Dict[str, torch.Tensor]:
+        zero_point: Tensor | None = None,
+        g_idx: torch.Tensor | None = None,
+        device: torch.device | None = None,
+        global_scale: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
         """
         Compresses a single uncompressed weight
 
@@ -146,8 +146,8 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
 
     def decompress_weight(
         self,
-        compressed_data: Dict[str, Tensor],
-        quantization_args: Optional[QuantizationArgs] = None,
+        compressed_data: dict[str, Tensor],
+        quantization_args: QuantizationArgs | None = None,
     ) -> torch.Tensor:
         """
         Decompresses a single compressed weight
@@ -190,7 +190,7 @@ class PackedQuantizationCompressor(BaseQuantizationCompressor):
 def pack_to_int32(
     value: torch.Tensor,
     num_bits: int,
-    packed_dim: Union[Literal[0], Literal[1]] = 1,
+    packed_dim: Literal[0, 1] = 1,
 ) -> torch.Tensor:
     """
     Packs a tensor of quantized weights stored in int8 into int32s with padding
@@ -254,7 +254,7 @@ def unpack_from_int32(
     value: torch.Tensor,
     num_bits: int,
     shape: torch.Size,
-    packed_dim: Union[Literal[0], Literal[1]] = 1,
+    packed_dim: Literal[0, 1] = 1,
 ) -> torch.Tensor:
     """
     Unpacks a tensor of packed int32 weights into individual int8s, maintaining the

{compressed_tensors-0.13.1a20260130 → compressed_tensors-0.13.1a20260203}/src/compressed_tensors/compressors/sparse_compressors/base.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import logging
-from typing import Dict, Generator, Optional, Set, Tuple
+from collections.abc import Generator
 
 from compressed_tensors.compressors.base import BaseCompressor
 from compressed_tensors.utils import (
@@ -65,10 +65,10 @@ class BaseSparseCompressor(BaseCompressor):
 
     def compress(
         self,
-        model_state: Dict[str, Tensor],
-        compression_targets: Optional[Set[str]] = None,
+        model_state: dict[str, Tensor],
+        compression_targets: set[str] | None = None,
         show_progress: bool = False,
-    ) -> Dict[str, Tensor]:
+    ) -> dict[str, Tensor]:
         """
         Compresses a dense state dict using bitmask compression
 
@@ -110,9 +110,9 @@ class BaseSparseCompressor(BaseCompressor):
         self,
         path_to_model_or_tensors: str,
         device: str = "cpu",
-        params_to_skip_load: Optional[Tuple] = None,
+        params_to_skip_load: tuple | None = None,
         **kwargs,
-    ) -> Generator[Tuple[str, Tensor], None, None]:
+    ) -> Generator[tuple[str, Tensor], None, None]:
         """
         Reads a bitmask compressed state dict located
         at path_to_model_or_tensors and returns a generator
@@ -157,8 +157,8 @@ class BaseSparseCompressor(BaseCompressor):
 
     def decompress_from_state_dict(
         self,
-        state_dict: Dict[str, Tensor],
-    ) -> Generator[Tuple[str, Dict[str, Tensor]], None, None]:
+        state_dict: dict[str, Tensor],
+    ) -> Generator[tuple[str, dict[str, Tensor]], None, None]:
         """
         Decompress the state dict of a module (or model)
 
@@ -185,7 +185,7 @@ class BaseSparseCompressor(BaseCompressor):
             yield ignored_param_path, ignored_param_value
 
     @staticmethod
-    def should_compress(name: str, expanded_targets: Optional[Set[str]] = None) -> bool:
+    def should_compress(name: str, expanded_targets: set[str] | None = None) -> bool:
         """
         Check if a parameter should be compressed.
         Currently, this only returns True for weight parameters.