Address mypy suggestions for type annotations in cablock.py, downsample.py, restormer.py and test_downsample_block.py.

Signed-off-by: Cano-Muniz, Santiago <[email protected]>

Author: phisanti

monai/networks/blocks/cablock.py

@@ -10,6 +10,8 @@
 # limitations under the License.
 from __future__ import annotations
 
+from typing import cast
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -70,7 +72,7 @@ def __init__(self, spatial_dims: int, dim: int, ffn_expansion_factor: float, bia
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.project_in(x)
         x1, x2 = self.dwconv(x).chunk(2, dim=1)
-        return self.project_out(F.gelu(x1) * x2)
+        return cast(torch.Tensor, self.project_out(F.gelu(x1) * x2))
 
 
 class CABlock(nn.Module):
@@ -141,7 +143,7 @@ def _normal_attention(self, q, k, v):
         attn = attn.softmax(dim=-1)
         return attn @ v
 
-    def forward(self, x) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass for MDTA attention.
         1. Apply depth-wise convolutions to Q, K, V
         2. Reshape Q, K, V for multi-head attention
@@ -177,4 +179,4 @@ def forward(self, x) -> torch.Tensor:
             **dict(zip(["h", "w"] if self.spatial_dims == 2 else ["d", "h", "w"], spatial_dims)),
         )
 
-        return self.project_out(out)
+        return cast(torch.Tensor, self.project_out(out))

monai/networks/blocks/downsample.py

@@ -92,7 +92,7 @@ def __init__(
         out_channels: int | None = None,
         scale_factor: Sequence[float] | float = 2,
         kernel_size: Sequence[float] | float | None = None,
-        mode: str = "conv",  # conv, convgroup, nontrainable, pixelunshuffle
+        mode: DownsampleMode | str = DownsampleMode.CONV,
         pre_conv: nn.Module | str | None = "default",
         post_conv: nn.Module | None = None,
         bias: bool = True,
@@ -101,11 +101,11 @@ def __init__(
         Downsamples data by `scale_factor`.
         Supported modes are:
 
-            - "conv": uses a strided convolution for learnable downsampling.
-            - "convgroup": uses a grouped strided convolution for efficient feature reduction.
-            - "maxpool": uses maxpooling for non-learnable downsampling.
-            - "avgpool": uses average pooling for non-learnable downsampling.
-            - "pixelunshuffle": uses :py:class:`monai.networks.blocks.SubpixelDownsample`.
+            - DownsampleMode.CONV: uses a strided convolution for learnable downsampling.
+            - DownsampleMode.CONVGROUP: uses a grouped strided convolution for efficient feature reduction.
+            - DownsampleMode.MAXPOOL: uses maxpooling for non-learnable downsampling.
+            - DownsampleMode.AVGPOOL: uses average pooling for non-learnable downsampling.
+            - DownsampleMode.PIXELUNSHUFFLE: uses :py:class:`monai.networks.blocks.SubpixelDownsample`.
 
         This operation will cause non-deterministic behavior when ``mode`` is ``DownsampleMode.NONTRAINABLE``.
         Please check the link below for more details:
@@ -120,7 +120,8 @@ def __init__(
             out_channels: number of channels of the output image. Defaults to `in_channels`.
             scale_factor: multiplier for spatial size reduction. Has to match input size if it is a tuple. Defaults to 2.
             kernel_size: kernel size used during convolutions. Defaults to `scale_factor`.
-            mode: {``"conv"``, ``"convgroup"``, ``"maxpool"``, ``"avgpool"``, ``"pixelunshuffle"``}. Defaults to ``"conv"``.
+            mode: {``DownsampleMode.CONV``, ``DownsampleMode.CONVGROUP``, ``DownsampleMode.MAXPOOL``, ``DownsampleMode.AVGPOOL``,
+                ``DownsampleMode.PIXELUNSHUFFLE``}. Defaults to ``DownsampleMode.CONV``.
             pre_conv: a conv block applied before downsampling. Defaults to "default".
                 When ``conv_block`` is ``"default"``, one reserved conv layer will be utilized.
                 Only used in the "maxpool", "avgpool" or "pixelunshuffle" modes.
@@ -134,7 +135,7 @@ def __init__(
 
         if not kernel_size:
             kernel_size_ = scale_factor_
-            padding = 0
+            padding = ensure_tuple_rep(0, spatial_dims)
         else:
             kernel_size_ = ensure_tuple_rep(kernel_size, spatial_dims)
             padding = tuple((k - 1) // 2 for k in kernel_size_)

monai/networks/nets/restormer.py

@@ -15,9 +15,10 @@
 
 from monai.networks.blocks.cablock import CABlock, FeedForward
 from monai.networks.blocks.convolutions import Convolution
-from monai.networks.blocks.downsample import DownSample, DownsampleMode
-from monai.networks.blocks.upsample import UpSample, UpsampleMode
+from monai.networks.blocks.downsample import DownSample
+from monai.networks.blocks.upsample import UpSample
 from monai.networks.layers.factories import Norm
+from monai.utils.enums import DownsampleMode, UpsampleMode
 
 
 class MDTATransformerBlock(nn.Module):
@@ -81,9 +82,9 @@ def __init__(self, spatial_dims: int, in_channels: int = 3, embed_dim: int = 48,
             conv_only=True,
         )
 
-
-def forward(self, x: torch.Tensor) -> torch.Tensor:
-    return super().forward(x)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = super().forward(x)
+        return x
 
 
 class Restormer(nn.Module):
@@ -290,7 +291,7 @@ def __init__(
             conv_only=True,
         )
 
-    def forward(self, x) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass of Restormer.
         Processes input through encoder-decoder architecture with skip connections.
         Args:

tests/networks/blocks/test_downsample_block.py

@@ -146,7 +146,8 @@ def __init__(self, n_feat: int):
 
             def forward(self, x: torch.Tensor) -> torch.Tensor:
                 x = self.conv(x)
-                return self.pixelunshuffle(x)
+                x = self.pixelunshuffle(x)
+                return x
 
         n_feat = 2
         x = torch.randn(1, n_feat, 64, 64)