[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot] · pre-commit-ci[bot] · commit ff7ed550592b · 2025-09-08T16:53:19.000Z
for more information, see https://pre-commit.ci
diff --git a/src/torchgmm/base/data/collation.py b/src/torchgmm/base/data/collation.py
@@ -1,9 +1,7 @@
-from typing import Tuple
-
 import torch
 
 
-def collate_tuple(batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
+def collate_tuple(batch: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
     """
     Collate a tuple of batch items by returning the input tuple.
 
@@ -13,7 +11,7 @@ def collate_tuple(batch: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
     return batch
 
 
-def collate_tensor(batch: Tuple[torch.Tensor, ...]) -> torch.Tensor:
+def collate_tensor(batch: tuple[torch.Tensor, ...]) -> torch.Tensor:
     """
     Collates a tuple of batch items into the first tensor.
 
diff --git a/src/torchgmm/base/data/loader.py b/src/torchgmm/base/data/loader.py
@@ -1,4 +1,5 @@
-from typing import Any, Iterator, TypeVar
+from collections.abc import Iterator
+from typing import Any, TypeVar
 
 try:
     from pytorch_lightning.overrides.distributed import _IndexBatchSamplerWrapper as IndexBatchSamplerWrapper
diff --git a/src/torchgmm/base/data/sampler.py b/src/torchgmm/base/data/sampler.py
@@ -1,5 +1,5 @@
 import math
-from typing import Iterator
+from collections.abc import Iterator
 
 from torch.utils.data import Sampler
 from torch.utils.data.sampler import SequentialSampler
diff --git a/src/torchgmm/base/data/types.py b/src/torchgmm/base/data/types.py
@@ -1,4 +1,5 @@
-from typing import Sequence, TypeVar, Union
+from collections.abc import Sequence
+from typing import TypeVar, Union
 
 import numpy as np
 import numpy.typing as npt
diff --git a/src/torchgmm/base/nn/_protocols.py b/src/torchgmm/base/nn/_protocols.py
@@ -1,5 +1,7 @@
 # pylint: disable=missing-class-docstring,missing-function-docstring
-from typing import Generic, Iterator, OrderedDict, Protocol, Tuple, Type, TypeVar
+from collections import OrderedDict
+from collections.abc import Iterator
+from typing import Generic, Protocol, TypeVar
 
 import torch
 from torch import nn
@@ -15,13 +17,13 @@ class ConfigurableModule(Protocol, Generic[C_co]):
     def config(self) -> C_co: ...
 
     @classmethod
-    def load(cls: Type[M], path: PathType) -> M: ...
+    def load(cls: type[M], path: PathType) -> M: ...
 
     def save(self, path: PathType, compile_model: bool = False) -> None: ...
 
     def save_config(self, path: PathType) -> None: ...
 
-    def named_children(self) -> Iterator[Tuple[str, nn.Module]]: ...
+    def named_children(self) -> Iterator[tuple[str, nn.Module]]: ...
 
     def state_dict(self) -> OrderedDict[str, torch.Tensor]: ...
 
diff --git a/src/torchgmm/base/utils/generics.py b/src/torchgmm/base/utils/generics.py
@@ -1,7 +1,7 @@
-from typing import Any, Type, get_args, get_origin
+from typing import Any, get_args, get_origin
 
 
-def get_generic_type(cls: Type[Any], origin: Type[Any], index: int = 0) -> Type[Any]:
+def get_generic_type(cls: type[Any], origin: type[Any], index: int = 0) -> type[Any]:
     """
     Returns the ``index``-th generic type of the superclass ``origin``.
 
diff --git a/src/torchgmm/bayes/core/_jit.py b/src/torchgmm/bayes/core/_jit.py
@@ -14,7 +14,7 @@ def jit_log_normal(
         # Precision shape is `[num_components, dim, dim]`.
         log_prob = x.new_empty((x.size(0), means.size(0)))
         # We loop here to not blow up the size of intermediate matrices
-        for k, (mu, prec_chol) in enumerate(zip(means, precisions_cholesky)):
+        for k, (mu, prec_chol) in enumerate(zip(means, precisions_cholesky, strict=False)):
             inner = x.matmul(prec_chol) - mu.matmul(prec_chol)
             log_prob[:, k] = inner.square().sum(1)
     elif covariance_type == "tied":
diff --git a/src/torchgmm/bayes/gmm/estimator.py b/src/torchgmm/bayes/gmm/estimator.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, List, Tuple, cast
+from typing import Any, cast
 
 import numpy as np
 import torch
@@ -272,7 +272,7 @@ def score_samples(self, data: TensorLike) -> torch.Tensor:
             collate_fn=collate_tensor,
         )
         result = self.trainer().predict(GaussianMixtureLightningModule(self.model_), loader)
-        return torch.stack([x[1] for x in cast(List[Tuple[torch.Tensor, torch.Tensor]], result)])
+        return torch.stack([x[1] for x in cast(list[tuple[torch.Tensor, torch.Tensor]], result)])
 
     def predict(self, data: TensorLike) -> torch.Tensor:
         """
@@ -321,4 +321,4 @@ def predict_proba(self, data: TensorLike) -> torch.Tensor:
             collate_fn=collate_tensor,
         )
         result = self.trainer().predict(GaussianMixtureLightningModule(self.model_), loader)
-        return torch.cat([x[0] for x in cast(List[Tuple[torch.Tensor, torch.Tensor]], result)])
+        return torch.cat([x[0] for x in cast(list[tuple[torch.Tensor, torch.Tensor]], result)])
diff --git a/src/torchgmm/bayes/gmm/metrics.py b/src/torchgmm/bayes/gmm/metrics.py
@@ -1,4 +1,5 @@
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any
 
 import torch
 from torchmetrics import Metric
@@ -17,7 +18,7 @@ def __init__(
         self,
         num_components: int,
         *,
-        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
+        dist_sync_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
@@ -44,7 +45,7 @@ def __init__(
         num_components: int,
         num_features: int,
         *,
-        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
+        dist_sync_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
@@ -78,7 +79,7 @@ def __init__(
         covariance_type: CovarianceType,
         reg: float,
         *,
-        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
+        dist_sync_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
diff --git a/src/torchgmm/bayes/gmm/model.py b/src/torchgmm/bayes/gmm/model.py
@@ -1,5 +1,4 @@
 from dataclasses import dataclass
-from typing import Tuple
 
 import numpy as np
 import torch
@@ -90,7 +89,7 @@ def reset_parameters(self) -> None:
         if self.covariance_type in ("full", "tied"):
             self.precisions_cholesky.tril_()
 
-    def forward(self, data: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def forward(self, data: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Computes the log-probability of observing each of the provided datapoints for each of the
         GMM's components.
diff --git a/src/torchgmm/clustering/kmeans/estimator.py b/src/torchgmm/clustering/kmeans/estimator.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, List, cast
+from typing import Any, cast
 
 import numpy as np
 import torch
@@ -188,7 +188,7 @@ def predict(self, data: TensorLike) -> torch.Tensor:
             collate_fn=collate_tensor,
         )
         result = self.trainer().predict(KMeansLightningModule(self.model_, predict_target="assignments"), loader)
-        return torch.cat(cast(List[torch.Tensor], result))
+        return torch.cat(cast(list[torch.Tensor], result))
 
     def score(self, data: TensorLike) -> float:
         """
@@ -236,7 +236,7 @@ def score_samples(self, data: TensorLike) -> torch.Tensor:
             collate_fn=collate_tensor,
         )
         result = self.trainer().predict(KMeansLightningModule(self.model_, predict_target="inertias"), loader)
-        return torch.cat(cast(List[torch.Tensor], result))
+        return torch.cat(cast(list[torch.Tensor], result))
 
     def transform(self, data: TensorLike) -> torch.Tensor:
         """
@@ -262,4 +262,4 @@ def transform(self, data: TensorLike) -> torch.Tensor:
             collate_fn=collate_tensor,
         )
         result = self.trainer().predict(KMeansLightningModule(self.model_, predict_target="distances"), loader)
-        return torch.cat(cast(List[torch.Tensor], result))
+        return torch.cat(cast(list[torch.Tensor], result))
diff --git a/src/torchgmm/clustering/kmeans/lightning_module.py b/src/torchgmm/clustering/kmeans/lightning_module.py
@@ -1,6 +1,6 @@
 # pylint: disable=abstract-method
 import math
-from typing import List, Literal
+from typing import Literal
 
 import pytorch_lightning as pl
 import torch
@@ -56,7 +56,7 @@ def __init__(
         # Initialize metrics
         self.metric_inertia = MeanMetric()
 
-    def configure_callbacks(self) -> List[pl.Callback]:
+    def configure_callbacks(self) -> list[pl.Callback]:
         if self.convergence_tolerance == 0:
             return []
         early_stopping = EarlyStopping(
diff --git a/src/torchgmm/clustering/kmeans/metrics.py b/src/torchgmm/clustering/kmeans/metrics.py
@@ -1,5 +1,6 @@
 import random
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any
 
 import torch
 from torchmetrics import Metric
@@ -17,7 +18,7 @@ def __init__(
         num_clusters: int,
         num_features: int,
         *,
-        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
+        dist_sync_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
@@ -61,7 +62,7 @@ def __init__(
         num_choices: int,
         num_features: int,
         *,
-        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
+        dist_sync_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
@@ -123,7 +124,7 @@ def __init__(
         num_choices: int,
         num_features: int,
         *,
-        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
+        dist_sync_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
@@ -173,7 +174,7 @@ class BatchSummer(Metric):
 
     full_state_update = True
 
-    def __init__(self, num_values: int, *, dist_sync_fn: Optional[Callable[[Any], Any]] = None):
+    def __init__(self, num_values: int, *, dist_sync_fn: Callable[[Any], Any] | None = None):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
         self.sums: torch.Tensor
@@ -198,7 +199,7 @@ def __init__(
         num_values: int,
         for_variance: bool,
         *,
-        dist_sync_fn: Optional[Callable[[Any], Any]] = None,
+        dist_sync_fn: Callable[[Any], Any] | None = None,
     ):
         super().__init__(dist_sync_fn=dist_sync_fn)  # type: ignore
 
diff --git a/src/torchgmm/clustering/kmeans/model.py b/src/torchgmm/clustering/kmeans/model.py
@@ -1,5 +1,4 @@
 from dataclasses import dataclass
-from typing import Tuple
 
 import torch
 from torch import jit, nn
@@ -52,7 +51,7 @@ def reset_parameters(self) -> None:
         """
         nn.init.normal_(self.centroids)
 
-    def forward(self, data: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def forward(self, data: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
         Computes the distance of each datapoint to each centroid as well as the "inertia", the
         squared distance of each datapoint to its closest centroid.
diff --git a/src/torchgmm/utils/lightning_module.py b/src/torchgmm/utils/lightning_module.py
@@ -1,5 +1,4 @@
 from abc import ABC, abstractmethod
-from typing import List
 
 import pytorch_lightning as pl
 import torch
@@ -32,7 +31,7 @@ def on_train_epoch_end(self) -> None:
             self.nonparametric_training_epoch_end()
     else:
 
-        def training_epoch_end(self, outputs: List[torch.Tensor]) -> None:
+        def training_epoch_end(self, outputs: list[torch.Tensor]) -> None:
             """Training epoch end hook for PyTorch Lightning < 2.0.0."""
             self.nonparametric_training_epoch_end()
 
diff --git a/tests/_data/gmm.py b/tests/_data/gmm.py
@@ -1,5 +1,4 @@
 # pylint: disable=missing-function-docstring
-from typing import Tuple
 
 import torch
 
@@ -9,7 +8,7 @@
 
 def sample_gmm(
     num_datapoints: int, num_features: int, num_components: int, covariance_type: CovarianceType
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     config = GaussianMixtureModelConfig(num_components, num_features, covariance_type)
     model = GaussianMixtureModel(config)
 
diff --git a/tests/_data/normal.py b/tests/_data/normal.py
@@ -1,28 +1,27 @@
 # pylint: disable=missing-function-docstring
-from typing import List
 
 import torch
 
 
-def sample_data(counts: List[int], dims: List[int]) -> List[torch.Tensor]:
-    return [torch.randn(count, dim) for count, dim in zip(counts, dims)]
+def sample_data(counts: list[int], dims: list[int]) -> list[torch.Tensor]:
+    return [torch.randn(count, dim) for count, dim in zip(counts, dims, strict=False)]
 
 
-def sample_means(counts: List[int], dims: List[int]) -> List[torch.Tensor]:
-    return [torch.randn(count, dim) for count, dim in zip(counts, dims)]
+def sample_means(counts: list[int], dims: list[int]) -> list[torch.Tensor]:
+    return [torch.randn(count, dim) for count, dim in zip(counts, dims, strict=False)]
 
 
-def sample_spherical_covars(counts: List[int]) -> List[torch.Tensor]:
+def sample_spherical_covars(counts: list[int]) -> list[torch.Tensor]:
     return [torch.rand(count) for count in counts]
 
 
-def sample_diag_covars(counts: List[int], dims: List[int]) -> List[torch.Tensor]:
-    return [torch.rand(count, dim).squeeze() for count, dim in zip(counts, dims)]
+def sample_diag_covars(counts: list[int], dims: list[int]) -> list[torch.Tensor]:
+    return [torch.rand(count, dim).squeeze() for count, dim in zip(counts, dims, strict=False)]
 
 
-def sample_full_covars(counts: List[int], dims: List[int]) -> List[torch.Tensor]:
+def sample_full_covars(counts: list[int], dims: list[int]) -> list[torch.Tensor]:
     result = []
-    for count, dim in zip(counts, dims):
+    for count, dim in zip(counts, dims, strict=False):
         A = torch.rand(count, dim * 10, dim)
         result.append(A.permute(0, 2, 1).bmm(A).squeeze())
     return result
diff --git a/tests/bayes/core/test_normal.py b/tests/bayes/core/test_normal.py
@@ -161,6 +161,7 @@ def test_cholesky_logdet_tied(covars: torch.Tensor):
         sample_data([10, 50, 100], [3, 50, 100]),
         sample_means([70, 5, 200], [3, 50, 100]),
         sample_spherical_covars([70, 5, 200]),
+        strict=False,
     ),
 )
 def test_log_normal_spherical(x: torch.Tensor, means: torch.Tensor, covars: torch.Tensor):
@@ -176,6 +177,7 @@ def test_log_normal_spherical(x: torch.Tensor, means: torch.Tensor, covars: torc
         sample_data([10, 50, 100], [3, 50, 100]),
         sample_means([70, 5, 200], [3, 50, 100]),
         sample_diag_covars([70, 5, 200], [3, 50, 100]),
+        strict=False,
     ),
 )
 def test_log_normal_diag(x: torch.Tensor, means: torch.Tensor, covars: torch.Tensor):
@@ -191,6 +193,7 @@ def test_log_normal_diag(x: torch.Tensor, means: torch.Tensor, covars: torch.Ten
         sample_data([10, 50, 100], [3, 50, 100]),
         sample_means([70, 5, 200], [3, 50, 100]),
         sample_full_covars([70, 5, 200], [3, 50, 100]),
+        strict=False,
     ),
 )
 def test_log_normal_full(x: torch.Tensor, means: torch.Tensor, covars: torch.Tensor):
@@ -205,6 +208,7 @@ def test_log_normal_full(x: torch.Tensor, means: torch.Tensor, covars: torch.Ten
         sample_data([10, 50, 100], [3, 50, 100]),
         sample_means([70, 5, 200], [3, 50, 100]),
         sample_full_covars([1, 1, 1], [3, 50, 100]),
+        strict=False,
     ),
 )
 def test_log_normal_tied(x: torch.Tensor, means: torch.Tensor, covars: torch.Tensor):
diff --git a/tests/bayes/gmm/benchmark_gmm_estimator.py b/tests/bayes/gmm/benchmark_gmm_estimator.py
@@ -1,5 +1,4 @@
 # pylint: disable=missing-function-docstring
-from typing import Optional
 
 import pytest
 import pytorch_lightning as pl
@@ -72,7 +71,7 @@ def test_torchgmm(
     num_features: int,
     num_components: int,
     covariance_type: CovarianceType,
-    batch_size: Optional[int],
+    batch_size: int | None,
 ):
     pl.seed_everything(0)
     data, means = sample_gmm(num_datapoints, num_features, num_components, covariance_type)
@@ -115,7 +114,7 @@ def test_torchgmm_gpu(
     num_features: int,
     num_components: int,
     covariance_type: CovarianceType,
-    batch_size: Optional[int],
+    batch_size: int | None,
 ):
     # Initialize GPU
     torch.empty(1, device="cuda:0")
diff --git a/tests/bayes/gmm/test_gmm_estimator.py b/tests/bayes/gmm/test_gmm_estimator.py
diff --git a/tests/bayes/gmm/test_gmm_metrics.py b/tests/bayes/gmm/test_gmm_metrics.py
diff --git a/tests/clustering/kmeans/benchmark_kmeans_estimator.py b/tests/clustering/kmeans/benchmark_kmeans_estimator.py
diff --git a/tests/clustering/kmeans/test_kmeans_estimator.py b/tests/clustering/kmeans/test_kmeans_estimator.py
