enforce stricter type check

Signed-off-by: Hao Wu <[email protected]>

Author: skyw

.pre-commit-config.yaml

@@ -33,10 +33,11 @@ repos:
     - id: ruff-format
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.14.0
+    rev: v1.19.1
     hooks:
       - id: mypy
         exclude: ^docs|^tests|^benchmarks|^docker
+        additional_dependencies: ["torch"]
 
   - repo: local
     hooks:

emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py

@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Callable, Literal
+from typing import Callable, Literal, overload
 
 
 # TODO(@boxiangw): remove this once bump to python 3.12
@@ -181,6 +181,12 @@ def _apply_moment2_normalization(
         else:
             raise TypeError(f"Invalid second moment method: {self.moment2_method}")
 
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
     @torch.no_grad()  # type: ignore[misc]
     @override
     def step(self, closure: Callable[[], float] | None = None) -> float | None:

emerging_optimizers/orthogonalized_optimizers/mop.py

@@ -57,6 +57,7 @@ def __init__(
         def scaled_orthogonalize_fn(grad: torch.Tensor) -> torch.Tensor:
             orth_grad, _, S = polar_via_svd(grad, False)
 
+            scale_factor: float | torch.Tensor
             if scale_mode != "nuclear_norm":
                 scale_factor = muon.get_muon_scale_factor(grad.size(-2), grad.size(-1), mode=scale_mode)
             else:

emerging_optimizers/orthogonalized_optimizers/muon_utils.py

@@ -123,7 +123,7 @@ def newton_schulz(
     if tp_group is not None:
         X = distributed_normalize_p2(x, eps, tp_group)
     else:
-        X = torch.nn.functional.normalize(x, p=2, dim=(-2, -1), eps=eps)
+        X = torch.nn.functional.normalize(x, p=2, dim=(-2, -1), eps=eps)  # type: ignore[arg-type]
 
     if coefficient_type in _COEFFICIENT_SETS:
         coefficient_sets = _COEFFICIENT_SETS[coefficient_type]

emerging_optimizers/orthogonalized_optimizers/orthogonalized_optimizer.py

@@ -12,7 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable
+from typing import Any, Callable, overload
 
 
 # TODO(@boxiangw): remove this once bump to python 3.12
@@ -126,6 +126,12 @@ def __init__(
         super().__init__(params, default_args_dict)
         self.scaled_orthogonalize_fn = scaled_orthogonalize_fn
 
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
     @torch.no_grad()  # type: ignore[misc]
     @override
     def step(self, closure: Callable[[], float] | None = None) -> float | None:

emerging_optimizers/psgd/procrustes_step.py

@@ -65,15 +65,17 @@ def procrustes_step(
             # rotate Q as exp(a R) Q ~ (I + a R + a^2 R^2/2) Q with an optimal step size by line search
             # for 2nd order expansion, only expand exp(a R) to its 2nd term.
             # Q += _step_size * (RQ + 0.5 * _step_size * RRQ)
-            Q = torch.add(Q, torch.add(RQ, RRQ, alpha=0.5 * step_size), alpha=step_size)
+            Q = torch.add(Q, torch.add(RQ, RRQ, alpha=0.5 * step_size), alpha=step_size)  # type: ignore[call-overload]
         if order == 3:
             RRRQ = R @ RRQ
             tr_RRRQ = torch.trace(RRRQ)
             # for a 3rd order expansion, we take the larger root of the cubic.
             _step_size = (-tr_RRQ - torch.sqrt(tr_RRQ * tr_RRQ - 1.5 * tr_RQ * tr_RRRQ)) / (0.75 * tr_RRRQ)
             step_size = torch.clamp(_step_size, max=max_step_size)
             # Q += step_size * (RQ + 0.5 * step_size * (RRQ + 0.25 * step_size * RRRQ))
-            Q = torch.add(
-                Q, torch.add(RQ, torch.add(RRQ, RRRQ, alpha=0.25 * step_size), alpha=0.5 * step_size), alpha=step_size
+            Q = torch.add(  # type: ignore[call-overload]
+                Q,
+                torch.add(RQ, torch.add(RRQ, RRRQ, alpha=0.25 * step_size), alpha=0.5 * step_size),  # type: ignore[call-overload]
+                alpha=step_size,  # type: ignore[call-overload]
             )
     return Q

emerging_optimizers/psgd/psgd.py

@@ -13,7 +13,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
-from typing import Callable, override
+from typing import Callable, overload
+
+
+try:
+    from typing import override
+except ImportError:
+    from typing_extensions import override
 
 import torch
 from torch.optim.optimizer import ParamsT
@@ -85,6 +91,12 @@ def __init__(
         }
         super().__init__(params, defaults)
 
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
     @torch.no_grad()  # type: ignore[misc]
     @override
     def step(self, closure: Callable[[], float] | None = None) -> float | None:

emerging_optimizers/riemannian_optimizers/normalized_optimizer.py

@@ -12,7 +12,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Callable
+from typing import Callable, overload
+
+
+try:
+    from typing import override
+except ImportError:
+    from typing_extensions import override
 
 import torch
 from torch.optim.optimizer import Optimizer
@@ -65,8 +71,15 @@ def __init__(
         )
         super().__init__(params, defaults)
 
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
     @torch.no_grad()  # type: ignore[misc]
-    def step(self, closure: Callable[[], float] | None = None) -> float | None:
+    @override
+    def step(self, closure: None = None) -> float | None:
         """Performs a single optimization step.
 
         Args:
@@ -154,7 +167,14 @@ def __init__(
         )
         super().__init__(params, defaults)
 
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
     @torch.no_grad()  # type: ignore[misc]
+    @override
     def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Performs a single optimization step.
 

emerging_optimizers/soap/soap.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 from functools import partial
 from itertools import chain
-from typing import Callable
+from typing import Callable, overload
 
 
 # TODO(@boxiangw): remove this once bump to python 3.12
@@ -136,6 +136,12 @@ def __init__(
         }
         super().__init__(params, defaults)
 
+    @overload
+    def step(self, closure: None = ...) -> None: ...
+
+    @overload
+    def step(self, closure: Callable[[], float]) -> float: ...
+
     @torch.no_grad()  # type: ignore[misc]
     @override
     def step(self, closure: Callable[[], float] | None = None) -> float | None:

emerging_optimizers/utils/modules.py

@@ -16,6 +16,12 @@
 import math
 from typing import Any, Self
 
+
+try:
+    from typing import override
+except ImportError:
+    from typing_extensions import override
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -51,8 +57,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 
         assert self.padding_mode == "zeros", "Only zeros padding is supported"
 
-        self.weight: nn.Parameter[torch.Tensor]
-        self.bias: nn.Parameter[torch.Tensor] | None | str
+        self.weight: nn.Parameter
+        self.bias: nn.Parameter | None
 
         flat_weight_shape = [self.out_channels, math.prod(self.weight.shape[1:])]
         if self.bias is not None:
@@ -63,7 +69,6 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             flat_weight_buffer[..., -1].copy_(self.bias)
             del self.bias
             self.has_bias = True
-            self.bias = "dummy"  # Trick con1d.extra_repr() to not print bias=False
         else:
             flat_weight_buffer.copy_(self.weight.view(self.out_channels, -1))
             self.has_bias = False
@@ -98,6 +103,7 @@ def from_conv1d(cls, conv1d: nn.Conv1d) -> Self:
     def weight_shape(self) -> tuple[int, int, int]:
         return (self.out_channels, self.in_channels // self.groups, self.kernel_size[0])
 
+    @override
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.has_bias:
             weight = self.weight[..., :-1].view(self.weight_shape)
@@ -108,6 +114,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         return F.conv1d(x, weight, bias, self.stride, self.padding, self.dilation, self.groups)
 
+    @override
     def extra_repr(self) -> str:
         base_repr = super().extra_repr()
+        if self.has_bias:
+            base_repr += ", bias=True"
         return f"{base_repr}, flattened_param_shape={tuple(self.weight.shape)}"