fix: enable invalid assignment ty checks (#553)

* fix: enable invalid assignment ty checks

* fix: ruff errors

* chore: undo type checking for notebooks

* fix: type check fixes for ty v0.0.20

* fix: ruff & numpy docstring check incompatibility

Author: begumcig

pyproject.toml

@@ -3,7 +3,7 @@ exclude = ["test_*.py"]
 line-length = 121
 
 [tool.ruff.lint]
-ignore = ["ANN002", "ANN003", "ANN401", "C901", "D100", "D104", "D401", "D406", "C408", "PTH123", "N813"]
+ignore = ["ANN002", "ANN003", "ANN401", "C901", "D100", "D104", "D401", "D406", "D420", "C408", "PTH123", "N813"]
 select = ["A", "C", "CPY", "D", "E", "ERA", "F", "FIX", "I", "N", "SIM", "T20", "W", "FA", "PTH"]
 preview = true # enable preview features for copyright checking
 
@@ -22,7 +22,6 @@ author = "- Pruna AI GmbH"
 [tool.ty.rules]
 # Ignore rules that are stricter than mypy for transition period
 unresolved-import = "ignore"
-invalid-assignment = "ignore"  # mypy is more permissive with Any assignments
 call-non-callable = "ignore"   # mypy allows more dynamic method calls
 index-out-of-bounds = "ignore" # mypy is more permissive with tuple indexing
 unresolved-attribute = "ignore" # mypy is more permissive with module attributes
@@ -38,6 +37,7 @@ possibly-missing-attribute = "ignore"
 missing-argument = "ignore"
 unused-type-ignore-comment = "ignore"
 
+
 [tool.coverage.run]
 source = ["src/pruna"]
 

src/pruna/algorithms/base/registry.py

@@ -15,7 +15,7 @@
 import inspect
 import logging
 import pkgutil
-from typing import Any, Callable, Dict
+from typing import Any, Dict
 
 from pruna.algorithms.base.pruna_base import PrunaAlgorithmBase
 from pruna.algorithms.base.tags import AlgorithmTag
@@ -28,7 +28,7 @@ class AlgorithmRegistry:
     The registry is a dictionary that maps algorithm names to algorithm instances.
     """
 
-    _registry: Dict[str, Callable[..., Any]] = {}
+    _registry: Dict[str, PrunaAlgorithmBase] = {}
 
     @classmethod
     def discover_algorithms(cls, algorithms_pkg: Any) -> None:

src/pruna/algorithms/c_translate.py

@@ -26,6 +26,7 @@
 from ConfigSpace import OrdinalHyperparameter
 from transformers import (
     AutomaticSpeechRecognitionPipeline,
+    PretrainedConfig,
     WhisperConfig,
 )
 from transformers.modeling_utils import PreTrainedModel
@@ -63,7 +64,7 @@ class CTranslate(PrunaAlgorithmBase):
     """
 
     algorithm_name: str = "c_translate"
-    group_tags: list[str] = [tags.COMPILER]
+    group_tags: list[tags] = [tags.COMPILER]
     save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.save_before_apply
     references = {"GitHub": "https://github.com/OpenNMT/CTranslate2"}
     tokenizer_required: bool = True
@@ -345,6 +346,7 @@ def __init__(self, generator: PreTrainedModel, output_dir: str, tokenizer: PreTr
         self.output_dir = output_dir
         self.task = "generation"
         self.tokenizer = tokenizer
+        self.config: PretrainedConfig | None = None
 
     def __getattr__(self, name: str) -> Any:
         """
@@ -416,6 +418,7 @@ def __init__(self, translator: PreTrainedModel, output_dir: str, tokenizer: PreT
         self.output_dir = output_dir
         self.task = "translation"
         self.tokenizer = tokenizer
+        self.config: PretrainedConfig | None = None
 
     def __getattr__(self, name: str) -> Any:
         """
@@ -499,6 +502,7 @@ def __init__(self, whisper: Whisper, output_dir: str, processor: ProcessorMixin)
         self.processor = processor
         self.language = None
         self.prompt = None
+        self.config: PretrainedConfig | None = None
 
     def __getattr__(self, name: str) -> Any:
         """

src/pruna/algorithms/deepcache.py

@@ -33,7 +33,7 @@ class DeepCache(PrunaAlgorithmBase):
     """
 
     algorithm_name: str = "deepcache"
-    group_tags: list[str] = [tags.CACHER]
+    group_tags: list[tags] = [tags.CACHER]
     save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.reapply
     references: dict[str, str] = {
         "GitHub": "https://github.com/horseee/DeepCache",

src/pruna/algorithms/fastercache.py

@@ -47,7 +47,7 @@ class FasterCache(PrunaAlgorithmBase):
     """
 
     algorithm_name: str = "fastercache"
-    group_tags: list[str] = [tags.CACHER]
+    group_tags: list[tags] = [tags.CACHER]
     save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.reapply
     references: dict[str, str] = {
         "GitHub": "https://github.com/Vchitect/FasterCache",

src/pruna/algorithms/flash_attn3.py

@@ -41,7 +41,7 @@ class FlashAttn3(PrunaAlgorithmBase):
     """
 
     algorithm_name: str = "flash_attn3"
-    group_tags: list[str] = [tags.KERNEL]
+    group_tags: list[tags] = [tags.KERNEL]
     save_fn = SAVE_FUNCTIONS.reapply
     references: dict[str, str] = {
         "GitHub": "https://github.com/Dao-AILab/flash-attention",

src/pruna/algorithms/fora.py

@@ -37,7 +37,7 @@ class FORA(PrunaAlgorithmBase):
     """
 
     algorithm_name: str = "fora"
-    group_tags: list[str] = [tags.CACHER]
+    group_tags: list[tags] = [tags.CACHER]
     save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.reapply
     references: dict[str, str] = {"Paper": "https://arxiv.org/abs/2407.01425"}
     tokenizer_required: bool = False
@@ -163,8 +163,8 @@ def __init__(self, pipe: Any, interval: int, start_step: int, backbone_calls_per
         self.single_stream_blocks_forward: Dict[int, Callable] = {}
 
         # Use seperate caches for the two different transformer block types
-        self.double_stream_blocks_cache: Dict[int, Tuple[Any, Any]] = {}
-        self.single_stream_blocks_cache: Dict[int, Any] = {}
+        self.double_stream_blocks_cache: Dict[Tuple[int, int], Tuple[Any, Any]] = {}
+        self.single_stream_blocks_cache: Dict[Tuple[int, int], Any] = {}
 
     def get_cache_schedule(self, num_steps: int) -> list[int]:
         """

src/pruna/algorithms/global_utils/recovery/finetuners/text_to_image_distiller.py

@@ -14,30 +14,24 @@
 
 from __future__ import annotations
 
+import contextlib
 import functools
+import pathlib
 import random
 from typing import Any, List, Literal
 
 import pytorch_lightning as pl
 import torch
-from diffusers.optimization import get_scheduler
-from diffusers.utils import BaseOutput
-from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
-from pytorch_lightning.utilities.seed import isolate_rng
-
-try:
-    from bitsandbytes.optim import AdamW8bit  # type: ignore[import-untyped]
-except ImportError:
-    AdamW8bit = None
-
-import pathlib
-
 from ConfigSpace import (
     CategoricalHyperparameter,
     Constant,
     UniformFloatHyperparameter,
     UniformIntegerHyperparameter,
 )
+from diffusers.optimization import get_scheduler
+from diffusers.utils import BaseOutput
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+from pytorch_lightning.utilities.seed import isolate_rng
 
 from pruna.algorithms.global_utils.recovery.finetuners import PrunaFinetuner
 from pruna.algorithms.global_utils.recovery.finetuners.diffusers import utils
@@ -56,6 +50,10 @@
 from pruna.engine.utils import get_device, get_device_type
 from pruna.logging.logger import pruna_logger
 
+AdamW8bit: type[Any] | None = None
+with contextlib.suppress(ImportError):
+    from bitsandbytes.optim import AdamW8bit
+
 
 class TextToImageDistiller(PrunaFinetuner):
     """Distiller for text-to-image models."""

src/pruna/algorithms/global_utils/recovery/finetuners/text_to_image_finetuner.py

@@ -14,25 +14,20 @@
 
 from __future__ import annotations
 
+import contextlib
 from pathlib import Path
 from typing import Any, List, Literal, Tuple
 
 import pytorch_lightning as pl
 import torch
-from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
-from pytorch_lightning.utilities.seed import isolate_rng
-
-try:
-    from bitsandbytes.optim import AdamW8bit  # type: ignore[import-untyped]
-except ImportError:
-    AdamW8bit = None
-
 from ConfigSpace import (
     CategoricalHyperparameter,
     Constant,
     UniformFloatHyperparameter,
     UniformIntegerHyperparameter,
 )
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+from pytorch_lightning.utilities.seed import isolate_rng
 
 from pruna.algorithms.global_utils.recovery.finetuners import PrunaFinetuner
 from pruna.algorithms.global_utils.recovery.finetuners.diffusers import (
@@ -49,6 +44,10 @@
 from pruna.config.smash_config import SmashConfigPrefixWrapper
 from pruna.logging.logger import pruna_logger
 
+AdamW8bit: type[Any] | None = None
+with contextlib.suppress(ImportError):
+    from bitsandbytes.optim import AdamW8bit
+
 
 class TextToImageFinetuner(PrunaFinetuner):
     """Finetuner for text-to-image models."""

src/pruna/algorithms/gptq_model.py

@@ -38,7 +38,7 @@ class GPTQ(PrunaAlgorithmBase):
     """
 
     algorithm_name: str = "gptq"
-    group_tags: list[str] = [tags.QUANTIZER]
+    group_tags: list[tags] = [tags.QUANTIZER]
     references: dict[str, str] = {"GitHub": "https://github.com/ModelCloud/GPTQModel"}
     save_fn: None = None
     tokenizer_required: bool = True