fix: types

Author: walln

.github/workflows/test.yaml

@@ -27,7 +27,7 @@ jobs:
               run: uvx ruff format --check
 
             - name: Typecheck
-              run: uvx mypy src --enable-incomplete-feature=NewGenericSyntax
+              run: uvx ty check
 
             - name: Test
               run: uv run pytest tests

.pre-commit-config.yaml

@@ -21,7 +21,7 @@ repos:
         stages: [commit-msg]
       - id: typecheck
         name: typecheck
-        entry: uvx mypy src/scratch/
+        entry: uvx ty check
         language: system
         types: [python]
         pass_filenames: false

pyproject.toml

@@ -42,6 +42,8 @@ dev = [
     "ipykernel>=6.29.5",
     "jupyter>=1.1.1",
     "ruff>=0.8.4",
+    "ty>=0.0.1a10",
+
 ]
 
 [tool.uv]
@@ -62,6 +64,8 @@ packages = ["src/scratch"]
 target-version = "py312"
 include = ["src/**", "tests/**"]
 line-length = 88
+
+[tool.ruff.lint]
 select = [
     "E",      # pycodestyle
     "W",      # pycodestyle
@@ -79,10 +83,10 @@ select = [
 ]
 ignore = ["F722"]
 
-[tool.ruff.pydocstyle]
+[tool.ruff.lint.pydocstyle]
 convention = "google"
 
-[tool.ruff.isort]
+[tool.ruff.lint.isort]
 known-first-party = ["src"]
 
 [tool.pytest.ini_options]

src/scratch/datasets/utils.py

@@ -4,7 +4,8 @@
 import warnings
 from dataclasses import dataclass, field
 
-from transformers import AutoTokenizer, PreTrainedTokenizerBase
+from transformers import AutoTokenizer
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 
 def patch_datasets_warning():
@@ -32,16 +33,6 @@ def filter_specific_warning(warning):
                 frame = frame.f_back
         return False
 
-    # Register the custom filter
-    warnings.filterwarnings("ignore", category=UserWarning, module=r".*")
-    warnings.showwarning = (
-        lambda message, category, filename, lineno, file=None, line=None: None
-        if filter_specific_warning(
-            warnings.WarningMessage(message, category, filename, lineno)
-        )
-        else warnings.showwarning(message, category, filename, lineno)
-    )
-
 
 @dataclass
 class TokenizerMetadata:
@@ -62,7 +53,7 @@ class TokenizerMetadata:
     def from_tokenizer(cls, tokenizer: PreTrainedTokenizerBase, max_length: int):
         """Create metadata from a tokenizer instance."""
         vocab_size = tokenizer.vocab_size  # type: ignore
-        if not vocab_size:
+        if not vocab_size or not isinstance(vocab_size, int):
             raise ValueError("The tokenizer does not have a vocab size.")
         return cls(
             vocab_size=vocab_size,

src/scratch/deep_learning/layers/attention/rope.py

@@ -92,6 +92,10 @@ def precompute_theta_pos_freqs(dim: int, end: int, theta: float = 10000.0):
     """
     freqs = 1.0 / (theta ** (jnp.arange(0, dim, 2)[: (dim // 2)] / dim))
     t = jnp.arange(end, dtype=jnp.float32)
+
+    assert isinstance(t, jnp.ndarray)
+    assert isinstance(freqs, jnp.ndarray)
+
     freqs = jnp.outer(t, freqs)
     freqs_cis = jnp.exp(1j * freqs)  # Using Euler's formula to create complex numbers
     return freqs_cis

src/scratch/image_classification/cnn.py

@@ -112,5 +112,7 @@ def __call__(self, x):
     # And comment out the following line
     logger = None
 
-    trainer = ImageClassificationParallelTrainer(model, trainer_config, logger=logger)
+    trainer = ImageClassificationParallelTrainer[CNN](
+        model, trainer_config, logger=logger
+    )
     trainer.train_and_evaluate(dataset.train, dataset.test)

src/scratch/image_classification/resnet.py

@@ -321,5 +321,5 @@ def __call__(self, x):
     trainer_config = ImageClassificationParallelTrainerConfig(
         batch_size=batch_size, learning_rate=0.01, epochs=3
     )
-    trainer = ImageClassificationParallelTrainer(model, trainer_config)
+    trainer = ImageClassificationParallelTrainer[ResNet](model, trainer_config)
     trainer.train_and_evaluate(dataset.train, dataset.test)

src/scratch/image_classification/swin_transformer.py

@@ -741,5 +741,5 @@ def __call__(self, x: jnp.ndarray, train=True):
     trainer_config = ImageClassificationParallelTrainerConfig(
         batch_size=batch_size, epochs=5
     )
-    trainer = ImageClassificationParallelTrainer(model, trainer_config)
+    trainer = ImageClassificationParallelTrainer[SwinTransformer](model, trainer_config)
     trainer.train_and_evaluate(dataset.train, dataset.test)

src/scratch/image_classification/trainer.py

@@ -5,9 +5,8 @@
 devices.
 """
 
-from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TypeVar
+from typing import Protocol, TypeVar
 
 import jax
 import jax.numpy as jnp
@@ -26,6 +25,14 @@
 M = TypeVar("M", bound=nnx.Module)
 
 
+class CallableModule(Protocol):
+    """Protocol for callable modules."""
+
+    def __call__(self, *args, **kwargs) -> jnp.ndarray:
+        """Call the module."""
+        ...  # pragma: no cover
+
+
 @dataclass
 class ImageClassificationParallelTrainerConfig(SupervisedTrainerConfig):
     """Configuration for the ImageClassificationParallelTrainer."""
@@ -83,8 +90,8 @@ def train(
         def train_step(
             model: M, train_state: TrainState, inputs: jnp.ndarray, targets: jnp.ndarray
         ):
-            def loss_fn(model: Callable):
-                logits = model(inputs)
+            def loss_fn(model: nnx.Module):
+                logits = model(inputs)  # type: ignore
                 assert logits.shape == targets.shape
                 loss = optax.softmax_cross_entropy(logits=logits, labels=targets).mean()
                 return loss, logits

src/scratch/image_classification/vision_transformer.py

@@ -214,5 +214,7 @@ def img_to_patch(x: jnp.ndarray, patch_size: int):
     trainer_config = ImageClassificationParallelTrainerConfig(
         batch_size=batch_size, epochs=5
     )
-    trainer = ImageClassificationParallelTrainer(model, trainer_config)
+    trainer = ImageClassificationParallelTrainer[VisionTransformer](
+        model, trainer_config
+    )
     trainer.train_and_evaluate(dataset.train, dataset.test)