Merge branch 'dev' into precommit_autofixing

Author: ericspod

.github/workflows/codeql-analysis.yml

@@ -9,67 +9,110 @@
 # the `language` matrix defined below to confirm you have the correct set of
 # supported CodeQL languages.
 #
-name: "CodeQL"
+name: "CodeQL Advanced"
 
 on:
   push:
     branches: [ dev, main ]
   pull_request:
-    # The branches below must be a subset of the branches above
     branches: [ dev ]
   schedule:
-    - cron: '18 1 * * 0'
+    - cron: '0 2 * * 1'  # 2AM Monday
+
+env:
+  PYTHON_VER: '3.11'
+  PYTORCH_VER: '2.8.0'
+  BUILD_MONAI: 1
+  PIP_EXTRA_INDEX_URL: "https://download.pytorch.org/whl/cpu"  # forces CPU PyTorch installation, should be faster
 
 jobs:
   analyze:
-    name: Analyze
+    name: Analyze (${{ matrix.language }})
+    # Runner size impacts CodeQL analysis time. To learn more, please see:
+    #   - https://gh.io/recommended-hardware-resources-for-running-codeql
+    #   - https://gh.io/supported-runners-and-hardware-resources
+    #   - https://gh.io/using-larger-runners (GitHub.com only)
+    # Consider using larger runners or machines with greater resources for possible analysis time improvements.
     runs-on: ubuntu-latest
     permissions:
+      # required for all workflows
+      security-events: write
+
+      # required to fetch internal or private CodeQL packs
+      packages: read
+
+      # only required for workflows in private repositories
       actions: read
       contents: read
-      security-events: write
 
     strategy:
       fail-fast: false
       matrix:
-        language: [ 'cpp', 'python' ]
-        # CodeQL supports [ 'cpp', 'csharp', 'go', 'java', 'javascript', 'python', 'ruby' ]
-        # Learn more about CodeQL language support at https://git.io/codeql-language-support
-
+        include:
+        - language: actions
+          build-mode: none
+        - language: c-cpp
+          build-mode: none  # TODO: get Cpp building working, autobuild doesn't work and manual fails for inexplicable reasons.
+        - language: python
+          build-mode: none
+        # CodeQL supports the following values keywords for 'language': 'actions', 'c-cpp', 'csharp', 'go', 'java-kotlin', 'javascript-typescript', 'python', 'ruby', 'rust', 'swift'
+        # Use `c-cpp` to analyze code written in C, C++ or both
+        # Use 'java-kotlin' to analyze code written in Java, Kotlin or both
+        # Use 'javascript-typescript' to analyze code written in JavaScript, TypeScript or both
+        # To learn more about changing the languages that are analyzed or customizing the build mode for your analysis,
+        # see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/customizing-your-advanced-setup-for-code-scanning.
+        # If you are analyzing a compiled language, you can modify the 'build-mode' for that language to customize how
+        # your codebase is analyzed, see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/codeql-code-scanning-for-compiled-languages
     steps:
     - name: Checkout repository
       uses: actions/checkout@v7
 
+    # Add any setup steps before running the `github/codeql-action/init` action.
+    # This includes steps like installing compilers or runtimes (`actions/setup-node`
+    # or others). This is typically only required for manual builds.
+    # - name: Setup runtime (example)
+    #   uses: actions/setup-example@v1
+
     # Initializes the CodeQL tools for scanning.
     - name: Initialize CodeQL
       uses: github/codeql-action/init@v4
       with:
         languages: ${{ matrix.language }}
+        build-mode: ${{ matrix.build-mode }}
         # If you wish to specify custom queries, you can do so here or in a config file.
         # By default, queries listed here will override any specified in a config file.
         # Prefix the list here with "+" to use these queries and those in the config file.
-        # queries: ./path/to/local/query, your-org/your-repo/queries@main
 
-    # Autobuild attempts to build any compiled languages  (C/C++, C#, or Java).
-    # If this step fails, then you should remove it and run the build manually (see below)
-    # - name: Autobuild
-    #   uses: github/codeql-action/autobuild@v2
+        # For more details on CodeQL's query packs, refer to: https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs
+        # queries: security-extended,security-and-quality
 
+    # If the analyze step fails for one of the languages you are analyzing with
+    # "We were unable to automatically build your code", modify the matrix above
+    # to set the build mode to "manual" for that language. Then modify this step
+    # to build your code.
     # ℹ️ Command-line programs to run using the OS shell.
-    # 📚 https://git.io/JvXDl
-
-    # ✏️ If the Autobuild fails above, remove it and uncomment the following three lines
-    #    and modify them (or add more) to build your code if your project
-    #    uses a compiled language
+    # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun
+    - name: Set up Python ${{ env.PYTHON_VER }}
+      if: matrix.language == 'c-cpp' && matrix.build-mode == 'manual'
+      uses: actions/setup-python@v6
+      with:
+        python-version: ${{ env.PYTHON_VER }}
+        cache: 'pip'
 
-    - name: Build
+    - name: Run manual build steps
+      if: matrix.language == 'c-cpp' && matrix.build-mode == 'manual'
+      shell: bash
       run: |
         rm -rf /opt/hostedtoolcache/{node,go,Ruby,Java*}
         ls -al /opt/hostedtoolcache
-        rm -rf /usr/share/dotnet/
+        sudo rm -rf /usr/share/dotnet/
         python -m pip install -U pip wheel wheel-stub
-        python -m pip install .[all,testing]
-        BUILD_MONAI=1 ./runtests.sh --build
+        python -m pip install torch==${PYTORCH_VER} torchvision
+        python -m pip install --user --upgrade pip wheel
+        python monai/config/print_dependencies.py build-system | xargs pip install --no-build-isolation
+        python -m pip install --no-build-isolation .
 
     - name: Perform CodeQL Analysis
       uses: github/codeql-action/analyze@v4
+      with:
+        category: "/language:${{matrix.language}}"

monai/networks/nets/dints.py

@@ -36,20 +36,23 @@
 __all__ = ["DiNTS", "TopologyConstruction", "TopologyInstance", "TopologySearch"]
 
 
-@torch.jit.interface
-class CellInterface(torch.nn.Module):
-    """interface for torchscriptable Cell"""
+# TODO: added temporarily for PyTorch 2.14 warnings, remove when factoring out deprecated Torchscript components
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
 
-    def forward(self, x: torch.Tensor, weight: torch.Tensor | None) -> torch.Tensor:  # type: ignore
-        pass
+    @torch.jit.interface
+    class CellInterface(torch.nn.Module):
+        """interface for torchscriptable Cell"""
 
+        def forward(self, x: torch.Tensor, weight: torch.Tensor | None) -> torch.Tensor:  # type: ignore
+            pass
 
-@torch.jit.interface
-class StemInterface(torch.nn.Module):
-    """interface for torchscriptable Stem"""
+    @torch.jit.interface
+    class StemInterface(torch.nn.Module):
+        """interface for torchscriptable Stem"""
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
-        pass
+        def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
+            pass
 
 
 class StemTS(StemInterface):

monai/networks/nets/swin_unetr.py

@@ -98,6 +98,7 @@ def __init__(
         hyena_omega_0: float = 10.0,
         hyena_l_cache: int = 32,
         hyena_short_conv_fft_chunks: int = 0,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -149,6 +150,7 @@ def __init__(
             hyena_omega_0: SIREN frequency. Default 10.0 (stable).
             hyena_l_cache: SIREN coordinate-grid cache size per spatial dim.
             hyena_short_conv_fft_chunks: channel chunk size for the FFT short conv (0 = no chunking).
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
 
         Examples::
 
@@ -224,6 +226,7 @@ def __init__(
             hyena_omega_0=hyena_omega_0,
             hyena_l_cache=hyena_l_cache,
             hyena_short_conv_fft_chunks=hyena_short_conv_fft_chunks,
+            use_flash_attention=use_flash_attention,
         )
 
         self.encoder1 = UnetrBasicBlock(
@@ -513,6 +516,7 @@ def __init__(
         qkv_bias: bool = False,
         attn_drop: float = 0.0,
         proj_drop: float = 0.0,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -522,12 +526,17 @@ def __init__(
             qkv_bias: add a learnable bias to query, key, value.
             attn_drop: attention dropout rate.
             proj_drop: dropout rate of output.
+            use_flash_attention: if True, use ``torch.nn.functional.scaled_dot_product_attention`` for the
+                windowed attention. Equivalent to the default path but faster at inference; only used when
+                autograd is disabled (e.g. under ``torch.no_grad()`` or ``torch.inference_mode()``, not
+                ``eval()`` alone) and the module is not scripted.
         """
 
         super().__init__()
         self.dim = dim
         self.window_size = window_size
         self.num_heads = num_heads
+        self.use_flash_attention = use_flash_attention
         head_dim = dim // num_heads
         self.scale = head_dim**-0.5
         mesh_args = torch.meshgrid.__kwdefaults__
@@ -584,12 +593,26 @@ def forward(self, x, mask):
         b, n, c = x.shape
         qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, c // self.num_heads).permute(2, 0, 3, 1, 4)
         q, k, v = qkv[0], qkv[1], qkv[2]
-        q = q * self.scale
-        attn = q @ k.transpose(-2, -1)
         relative_position_bias = self.relative_position_bias_table[
             self.relative_position_index.clone()[:n, :n].reshape(-1)  # type: ignore[operator]
         ].reshape(n, n, -1)
         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+        if self.use_flash_attention and not torch.jit.is_scripting() and not torch.is_grad_enabled():
+            # additive bias combines the relative position bias and, for shifted windows, the attention mask
+            if mask is not None:
+                nw = mask.shape[0]
+                bias = relative_position_bias.view(1, 1, self.num_heads, n, n) + mask.reshape(1, nw, 1, n, n)
+                bias = bias.expand(b // nw, nw, self.num_heads, n, n).reshape(b, self.num_heads, n, n)
+            else:
+                bias = relative_position_bias.unsqueeze(0)
+            x = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=bias.to(q.dtype), dropout_p=0.0, scale=self.scale
+            )
+            x = x.transpose(1, 2).reshape(b, n, c)
+            return self.proj_drop(self.proj(x))
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
         attn = attn + relative_position_bias.unsqueeze(0)
         if mask is not None:
             nw = mask.shape[0]
@@ -628,6 +651,7 @@ def __init__(
         act_layer: str = "GELU",
         norm_layer: type[LayerNorm] = nn.LayerNorm,
         use_checkpoint: bool = False,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -643,6 +667,7 @@ def __init__(
             act_layer: activation layer.
             norm_layer: normalization layer.
             use_checkpoint: use gradient checkpointing for reduced memory usage.
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
         """
 
         super().__init__()
@@ -660,6 +685,7 @@ def __init__(
             qkv_bias=qkv_bias,
             attn_drop=attn_drop,
             proj_drop=drop,
+            use_flash_attention=use_flash_attention,
         )
 
         self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
@@ -924,6 +950,7 @@ def __init__(
         hyena_omega_0: float = 10.0,
         hyena_l_cache: int = 32,
         hyena_short_conv_fft_chunks: int = 0,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -946,6 +973,7 @@ def __init__(
                 hyena_use_chunked_fft, hyena_use_fft_short_conv, hyena_omega_0, hyena_l_cache,
                 hyena_short_conv_fft_chunks: forwarded to :class:`HyenaTransformerBlock`. See its
                 docstring for semantics.
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
         """
 
         super().__init__()
@@ -996,6 +1024,7 @@ def __init__(
                         drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
                         norm_layer=norm_layer,
                         use_checkpoint=use_checkpoint,
+                        use_flash_attention=use_flash_attention,
                     )
                     for i in range(depth)
                 ]
@@ -1079,6 +1108,7 @@ def __init__(
         hyena_omega_0: float = 10.0,
         hyena_l_cache: int = 32,
         hyena_short_conv_fft_chunks: int = 0,
+        use_flash_attention: bool = False,
     ) -> None:
         """
         Args:
@@ -1112,6 +1142,7 @@ def __init__(
                 hyena_use_chunked_fft, hyena_use_fft_short_conv, hyena_omega_0, hyena_l_cache,
                 hyena_short_conv_fft_chunks: HyenaND configuration. See
                 :class:`monai.networks.blocks.HyenaTransformerBlock` for semantics.
+            use_flash_attention: use flash attention (scaled dot product attention) at inference.
         """
 
         super().__init__()
@@ -1193,6 +1224,7 @@ def __init__(
                 hyena_omega_0=hyena_omega_0,
                 hyena_l_cache=hyena_l_cache,
                 hyena_short_conv_fft_chunks=hyena_short_conv_fft_chunks,
+                use_flash_attention=use_flash_attention,
             )
             if i_layer == 0:
                 self.layers1.append(layer)

monai/transforms/croppad/array.py

@@ -342,6 +342,24 @@ def compute_pad_width(self, spatial_shape: Sequence[int]) -> tuple[tuple[int, in
         return spatial_pad.compute_pad_width(spatial_shape)
 
 
+def _to_int_list(data: Sequence[int] | int | NdarrayOrTensor) -> list[int]:
+    """Coerce an ROI spec (scalar, sequence, tensor or ndarray) to a list of Python ints."""
+    if isinstance(data, (str, bytes)):
+        raise TypeError("ROI specs must be integers or sequences of integers, not strings.")
+    return [int(i) for i in ensure_tuple(data)]
+
+
+def _broadcast_int_pair(
+    a: Sequence[int] | int | NdarrayOrTensor, b: Sequence[int] | int | NdarrayOrTensor
+) -> tuple[list[int], list[int]]:
+    """Coerce a pair of ROI specs to two equal-length int lists, broadcasting a scalar to match."""
+    list_a, list_b = _to_int_list(a), _to_int_list(b)
+    n = max(len(list_a), len(list_b))
+    if len(list_a) not in (1, n) or len(list_b) not in (1, n):
+        raise ValueError(f"ROI specs must have matching lengths or be scalar, got {len(list_a)} and {len(list_b)}.")
+    return (list_a * n if len(list_a) == 1 else list_a), (list_b * n if len(list_b) == 1 else list_b)
+
+
 class Crop(InvertibleTransform, LazyTransform):
     """
     Perform crop operations on the input image.
@@ -379,31 +397,22 @@ def compute_slices(
             roi_slices: list of slices for each of the spatial dimensions.
 
         """
-        roi_start_t: torch.Tensor
-
         if roi_slices:
             if not all(s.step is None or s.step == 1 for s in roi_slices):
                 raise ValueError(f"only slice steps of 1/None are currently supported, got {roi_slices}.")
             return ensure_tuple(roi_slices)
         else:
             if roi_center is not None and roi_size is not None:
-                roi_center_t = convert_to_tensor(data=roi_center, dtype=torch.int16, wrap_sequence=True, device="cpu")
-                roi_size_t = convert_to_tensor(data=roi_size, dtype=torch.int16, wrap_sequence=True, device="cpu")
-                _zeros = torch.zeros_like(roi_center_t)
-                half = torch.divide(roi_size_t, 2, rounding_mode="floor")
-                roi_start_t = torch.maximum(roi_center_t - half, _zeros)
-                roi_end_t = torch.maximum(roi_start_t + roi_size_t, roi_start_t)
+                centers, sizes = _broadcast_int_pair(roi_center, roi_size)
+                starts = [max(c - s // 2, 0) for c, s in zip(centers, sizes)]
+                ends = [st + s for st, s in zip(starts, sizes)]
             else:
                 if roi_start is None or roi_end is None:
                     raise ValueError("please specify either roi_center, roi_size or roi_start, roi_end.")
-                roi_start_t = convert_to_tensor(data=roi_start, dtype=torch.int16, wrap_sequence=True)
-                roi_start_t = torch.maximum(roi_start_t, torch.zeros_like(roi_start_t))
-                roi_end_t = convert_to_tensor(data=roi_end, dtype=torch.int16, wrap_sequence=True)
-                roi_end_t = torch.maximum(roi_end_t, roi_start_t)
-            # convert to slices (accounting for 1d)
-            if roi_start_t.numel() == 1:
-                return ensure_tuple([slice(int(roi_start_t.item()), int(roi_end_t.item()))])
-            return ensure_tuple([slice(int(s), int(e)) for s, e in zip(roi_start_t.tolist(), roi_end_t.tolist())])
+                starts, ends = _broadcast_int_pair(roi_start, roi_end)
+                starts = [max(s, 0) for s in starts]
+            # clamp each end to its own start so no slice has negative width
+            return ensure_tuple(slice(s, max(e, s)) for s, e in zip(starts, ends))
 
     def __call__(  # type: ignore[override]
         self, img: torch.Tensor, slices: tuple[slice, ...], lazy: bool | None = None