dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/models/sam/modules/transformer.py

@@ -1,7 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
 import math
-from typing import Tuple, Type
 
 import torch
 from torch import Tensor, nn
@@ -10,12 +11,10 @@ from ultralytics.nn.modules import MLPBlock
 
 
 class TwoWayTransformer(nn.Module):
-    """
-    A Two-Way Transformer module for simultaneous attention to image and query points.
+    """A Two-Way Transformer module for simultaneous attention to image and query points.
 
-    This class implements a specialized transformer decoder that attends to an input image using queries with
-    supplied positional embeddings. It's useful for tasks like object detection, image segmentation, and point
-    cloud processing.
+    This class implements a specialized transformer decoder that attends to an input image using queries with supplied
+    positional embeddings. It's useful for tasks like object detection, image segmentation, and point cloud processing.
 
     Attributes:
         depth (int): Number of layers in the transformer.
@@ -27,7 +26,7 @@ class TwoWayTransformer(nn.Module):
         norm_final_attn (nn.LayerNorm): Layer normalization applied to final queries.
 
     Methods:
-        forward: Processes image and point embeddings through the transformer.
+        forward: Process image and point embeddings through the transformer.
 
     Examples:
         >>> transformer = TwoWayTransformer(depth=6, embedding_dim=256, num_heads=8, mlp_dim=2048)
@@ -44,19 +43,18 @@ class TwoWayTransformer(nn.Module):
         embedding_dim: int,
         num_heads: int,
         mlp_dim: int,
-        activation: Type[nn.Module] = nn.ReLU,
+        activation: type[nn.Module] = nn.ReLU,
         attention_downsample_rate: int = 2,
     ) -> None:
-        """
-        Initialize a Two-Way Transformer for simultaneous attention to image and query points.
+        """Initialize a Two-Way Transformer for simultaneous attention to image and query points.
 
         Args:
             depth (int): Number of layers in the transformer.
             embedding_dim (int): Channel dimension for input embeddings.
             num_heads (int): Number of heads for multihead attention. Must divide embedding_dim.
             mlp_dim (int): Internal channel dimension for the MLP block.
-            activation (Type[nn.Module]): Activation function to use in the MLP block.
-            attention_downsample_rate (int): Downsampling rate for attention mechanism.
+            activation (Type[nn.Module], optional): Activation function to use in the MLP block.
+            attention_downsample_rate (int, optional): Downsampling rate for attention mechanism.
         """
         super().__init__()
         self.depth = depth
@@ -82,21 +80,20 @@ class TwoWayTransformer(nn.Module):
 
     def forward(
         self,
-        image_embedding: Tensor,
-        image_pe: Tensor,
-        point_embedding: Tensor,
-    ) -> Tuple[Tensor, Tensor]:
-        """
-        Process image and point embeddings through the Two-Way Transformer.
+        image_embedding: torch.Tensor,
+        image_pe: torch.Tensor,
+        point_embedding: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Process image and point embeddings through the Two-Way Transformer.
 
         Args:
-            image_embedding (Tensor): Image to attend to, with shape (B, embedding_dim, H, W).
-            image_pe (Tensor): Positional encoding to add to the image, with same shape as image_embedding.
-            point_embedding (Tensor): Embedding to add to query points, with shape (B, N_points, embedding_dim).
+            image_embedding (torch.Tensor): Image to attend to, with shape (B, embedding_dim, H, W).
+            image_pe (torch.Tensor): Positional encoding to add to the image, with same shape as image_embedding.
+            point_embedding (torch.Tensor): Embedding to add to query points, with shape (B, N_points, embedding_dim).
 
         Returns:
-            queries (Tensor): Processed point embeddings with shape (B, N_points, embedding_dim).
-            keys (Tensor): Processed image embeddings with shape (B, H*W, embedding_dim).
+            queries (torch.Tensor): Processed point embeddings with shape (B, N_points, embedding_dim).
+            keys (torch.Tensor): Processed image embeddings with shape (B, H*W, embedding_dim).
         """
         # BxCxHxW -> BxHWxC == B x N_image_tokens x C
         image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
@@ -126,12 +123,11 @@ class TwoWayTransformer(nn.Module):
 
 
 class TwoWayAttentionBlock(nn.Module):
-    """
-    A two-way attention block for simultaneous attention to image and query points.
+    """A two-way attention block for simultaneous attention to image and query points.
 
     This class implements a specialized transformer block with four main layers: self-attention on sparse inputs,
-    cross-attention of sparse inputs to dense inputs, MLP block on sparse inputs, and cross-attention of dense
-    inputs to sparse inputs.
+    cross-attention of sparse inputs to dense inputs, MLP block on sparse inputs, and cross-attention of dense inputs to
+    sparse inputs.
 
     Attributes:
         self_attn (Attention): Self-attention layer for queries.
@@ -145,7 +141,7 @@ class TwoWayAttentionBlock(nn.Module):
         skip_first_layer_pe (bool): Whether to skip positional encoding in the first layer.
 
     Methods:
-        forward: Applies self-attention and cross-attention to queries and keys.
+        forward: Apply self-attention and cross-attention to queries and keys.
 
     Examples:
         >>> embedding_dim, num_heads = 256, 8
@@ -162,24 +158,23 @@ class TwoWayAttentionBlock(nn.Module):
         embedding_dim: int,
         num_heads: int,
         mlp_dim: int = 2048,
-        activation: Type[nn.Module] = nn.ReLU,
+        activation: type[nn.Module] = nn.ReLU,
         attention_downsample_rate: int = 2,
         skip_first_layer_pe: bool = False,
     ) -> None:
-        """
-        Initialize a TwoWayAttentionBlock for simultaneous attention to image and query points.
+        """Initialize a TwoWayAttentionBlock for simultaneous attention to image and query points.
 
         This block implements a specialized transformer layer with four main components: self-attention on sparse
-        inputs, cross-attention of sparse inputs to dense inputs, MLP block on sparse inputs, and cross-attention
-        of dense inputs to sparse inputs.
+        inputs, cross-attention of sparse inputs to dense inputs, MLP block on sparse inputs, and cross-attention of
+        dense inputs to sparse inputs.
 
         Args:
             embedding_dim (int): Channel dimension of the embeddings.
             num_heads (int): Number of attention heads in the attention layers.
-            mlp_dim (int): Hidden dimension of the MLP block.
-            activation (Type[nn.Module]): Activation function for the MLP block.
-            attention_downsample_rate (int): Downsampling rate for the attention mechanism.
-            skip_first_layer_pe (bool): Whether to skip positional encoding in the first layer.
+            mlp_dim (int, optional): Hidden dimension of the MLP block.
+            activation (Type[nn.Module], optional): Activation function for the MLP block.
+            attention_downsample_rate (int, optional): Downsampling rate for the attention mechanism.
+            skip_first_layer_pe (bool, optional): Whether to skip positional encoding in the first layer.
         """
         super().__init__()
         self.self_attn = Attention(embedding_dim, num_heads)
@@ -196,19 +191,20 @@ class TwoWayAttentionBlock(nn.Module):
 
         self.skip_first_layer_pe = skip_first_layer_pe
 
-    def forward(self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor) -> Tuple[Tensor, Tensor]:
-        """
-        Apply two-way attention to process query and key embeddings in a transformer block.
+    def forward(
+        self, queries: torch.Tensor, keys: torch.Tensor, query_pe: torch.Tensor, key_pe: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply two-way attention to process query and key embeddings in a transformer block.
 
         Args:
-            queries (Tensor): Query embeddings with shape (B, N_queries, embedding_dim).
-            keys (Tensor): Key embeddings with shape (B, N_keys, embedding_dim).
-            query_pe (Tensor): Positional encodings for queries with same shape as queries.
-            key_pe (Tensor): Positional encodings for keys with same shape as keys.
+            queries (torch.Tensor): Query embeddings with shape (B, N_queries, embedding_dim).
+            keys (torch.Tensor): Key embeddings with shape (B, N_keys, embedding_dim).
+            query_pe (torch.Tensor): Positional encodings for queries with same shape as queries.
+            key_pe (torch.Tensor): Positional encodings for keys with same shape as keys.
 
         Returns:
-            queries (Tensor): Processed query embeddings with shape (B, N_queries, embedding_dim).
-            keys (Tensor): Processed key embeddings with shape (B, N_keys, embedding_dim).
+            queries (torch.Tensor): Processed query embeddings with shape (B, N_queries, embedding_dim).
+            keys (torch.Tensor): Processed key embeddings with shape (B, N_keys, embedding_dim).
         """
         # Self attention block
         if self.skip_first_layer_pe:
@@ -242,11 +238,10 @@ class TwoWayAttentionBlock(nn.Module):
 
 
 class Attention(nn.Module):
-    """
-    An attention layer with downscaling capability for embedding size after projection.
+    """An attention layer with downscaling capability for embedding size after projection.
 
-    This class implements a multi-head attention mechanism with the option to downsample the internal
-    dimension of queries, keys, and values.
+    This class implements a multi-head attention mechanism with the option to downsample the internal dimension of
+    queries, keys, and values.
 
     Attributes:
         embedding_dim (int): Dimensionality of input embeddings.
@@ -259,9 +254,9 @@ class Attention(nn.Module):
         out_proj (nn.Linear): Linear projection for output.
 
     Methods:
-        _separate_heads: Separates input tensor into attention heads.
-        _recombine_heads: Recombines separated attention heads.
-        forward: Computes attention output for given query, key, and value tensors.
+        _separate_heads: Separate input tensor into attention heads.
+        _recombine_heads: Recombine separated attention heads.
+        forward: Compute attention output for given query, key, and value tensors.
 
     Examples:
         >>> attn = Attention(embedding_dim=256, num_heads=8, downsample_rate=2)
@@ -277,16 +272,15 @@ class Attention(nn.Module):
         embedding_dim: int,
         num_heads: int,
         downsample_rate: int = 1,
-        kv_in_dim: int = None,
+        kv_in_dim: int | None = None,
     ) -> None:
-        """
-        Initialize the Attention module with specified dimensions and settings.
+        """Initialize the Attention module with specified dimensions and settings.
 
         Args:
             embedding_dim (int): Dimensionality of input embeddings.
             num_heads (int): Number of attention heads.
-            downsample_rate (int): Factor by which internal dimensions are downsampled.
-            kv_in_dim (int | None): Dimensionality of key and value inputs. If None, uses embedding_dim.
+            downsample_rate (int, optional): Factor by which internal dimensions are downsampled.
+            kv_in_dim (int | None, optional): Dimensionality of key and value inputs. If None, uses embedding_dim.
 
         Raises:
             AssertionError: If num_heads does not evenly divide the internal dim (embedding_dim / downsample_rate).
@@ -304,7 +298,7 @@ class Attention(nn.Module):
         self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
 
     @staticmethod
-    def _separate_heads(x: Tensor, num_heads: int) -> Tensor:
+    def _separate_heads(x: torch.Tensor, num_heads: int) -> torch.Tensor:
         """Separate the input tensor into the specified number of attention heads."""
         b, n, c = x.shape
         x = x.reshape(b, n, num_heads, c // num_heads)
@@ -317,17 +311,16 @@ class Attention(nn.Module):
         x = x.transpose(1, 2)
         return x.reshape(b, n_tokens, n_heads * c_per_head)  # B x N_tokens x C
 
-    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
-        """
-        Apply multi-head attention to query, key, and value tensors with optional downsampling.
+    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+        """Apply multi-head attention to query, key, and value tensors with optional downsampling.
 
         Args:
-            q (Tensor): Query tensor with shape (B, N_q, embedding_dim).
-            k (Tensor): Key tensor with shape (B, N_k, embedding_dim).
-            v (Tensor): Value tensor with shape (B, N_k, embedding_dim).
+            q (torch.Tensor): Query tensor with shape (B, N_q, embedding_dim).
+            k (torch.Tensor): Key tensor with shape (B, N_k, embedding_dim).
+            v (torch.Tensor): Value tensor with shape (B, N_k, embedding_dim).
 
         Returns:
-            (Tensor): Output tensor after attention with shape (B, N_q, embedding_dim).
+            (torch.Tensor): Output tensor after attention with shape (B, N_q, embedding_dim).
         """
         # Input projections
         q = self.q_proj(q)

ultralytics/models/sam/modules/utils.py

@@ -1,23 +1,24 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-from typing import Tuple
+from __future__ import annotations
+
+from typing import Any
 
 import torch
 import torch.nn.functional as F
 
 
-def select_closest_cond_frames(frame_idx, cond_frame_outputs, max_cond_frame_num):
-    """
-    Select the closest conditioning frames to a given frame index.
+def select_closest_cond_frames(frame_idx: int, cond_frame_outputs: dict[int, Any], max_cond_frame_num: int):
+    """Select the closest conditioning frames to a given frame index.
 
     Args:
         frame_idx (int): Current frame index.
-        cond_frame_outputs (Dict[int, Any]): Dictionary of conditioning frame outputs keyed by frame indices.
+        cond_frame_outputs (dict[int, Any]): Dictionary of conditioning frame outputs keyed by frame indices.
         max_cond_frame_num (int): Maximum number of conditioning frames to select.
 
     Returns:
-        selected_outputs (Dict[int, Any]): Selected items from cond_frame_outputs.
-        unselected_outputs (Dict[int, Any]): Items not selected from cond_frame_outputs.
+        selected_outputs (dict[int, Any]): Selected items from cond_frame_outputs.
+        unselected_outputs (dict[int, Any]): Items not selected from cond_frame_outputs.
 
     Examples:
         >>> frame_idx = 5
@@ -59,14 +60,13 @@ def select_closest_cond_frames(frame_idx, cond_frame_outputs, max_cond_frame_num
     return selected_outputs, unselected_outputs
 
 
-def get_1d_sine_pe(pos_inds, dim, temperature=10000):
-    """
-    Generate 1D sinusoidal positional embeddings for given positions and dimensions.
+def get_1d_sine_pe(pos_inds: torch.Tensor, dim: int, temperature: float = 10000):
+    """Generate 1D sinusoidal positional embeddings for given positions and dimensions.
 
     Args:
         pos_inds (torch.Tensor): Position indices for which to generate embeddings.
         dim (int): Dimension of the positional embeddings. Should be an even number.
-        temperature (float): Scaling factor for the frequency of the sinusoidal functions.
+        temperature (float, optional): Scaling factor for the frequency of the sinusoidal functions.
 
     Returns:
         (torch.Tensor): Sinusoidal positional embeddings with shape (pos_inds.shape, dim).
@@ -78,7 +78,7 @@ def get_1d_sine_pe(pos_inds, dim, temperature=10000):
         torch.Size([4, 128])
     """
     pe_dim = dim // 2
-    dim_t = torch.arange(pe_dim, dtype=torch.float32, device=pos_inds.device)
+    dim_t = torch.arange(pe_dim, dtype=pos_inds.dtype, device=pos_inds.device)
     dim_t = temperature ** (2 * (dim_t // 2) / pe_dim)
 
     pos_embed = pos_inds.unsqueeze(-1) / dim_t
@@ -87,25 +87,21 @@ def get_1d_sine_pe(pos_inds, dim, temperature=10000):
 
 
 def init_t_xy(end_x: int, end_y: int):
-    """
-    Initialize 1D and 2D coordinate tensors for a grid of specified dimensions.
+    """Initialize 1D and 2D coordinate tensors for a grid of specified dimensions.
 
-    This function creates coordinate tensors for a grid with dimensions end_x × end_y. It generates a linear index tensor
-    and corresponding x and y coordinate tensors.
+    This function creates coordinate tensors for a grid with dimensions end_x × end_y. It generates a linear index
+    tensor and corresponding x and y coordinate tensors.
 
     Args:
         end_x (int): Width of the grid (number of columns).
         end_y (int): Height of the grid (number of rows).
 
     Returns:
-        t (torch.Tensor): Linear indices for each position in the grid, with shape (end_x * end_y).
         t_x (torch.Tensor): X-coordinates for each position, with shape (end_x * end_y).
         t_y (torch.Tensor): Y-coordinates for each position, with shape (end_x * end_y).
 
     Examples:
-        >>> t, t_x, t_y = init_t_xy(3, 2)
-        >>> print(t)
-        tensor([0., 1., 2., 3., 4., 5.])
+        >>> t_x, t_y = init_t_xy(3, 2)
         >>> print(t_x)
         tensor([0., 1., 2., 0., 1., 2.])
         >>> print(t_y)
@@ -118,11 +114,10 @@ def init_t_xy(end_x: int, end_y: int):
 
 
 def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 10000.0):
-    """
-    Compute axial complex exponential positional encodings for 2D spatial positions in a grid.
+    """Compute axial complex exponential positional encodings for 2D spatial positions in a grid.
 
-    This function generates complex exponential positional encodings for a 2D grid of spatial positions,
-    using separate frequency components for the x and y dimensions.
+    This function generates complex exponential positional encodings for a 2D grid of spatial positions, using separate
+    frequency components for the x and y dimensions.
 
     Args:
         dim (int): Dimension of the positional encoding.
@@ -131,18 +126,13 @@ def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 10000.0):
         theta (float, optional): Scaling factor for frequency computation.
 
     Returns:
-        freqs_cis_x (torch.Tensor): Complex exponential positional encodings for x-dimension with shape
-            (end_x*end_y, dim//4).
-        freqs_cis_y (torch.Tensor): Complex exponential positional encodings for y-dimension with shape
-            (end_x*end_y, dim//4).
+        (torch.Tensor): Complex exponential positional encodings with shape (end_x*end_y, dim//2).
 
     Examples:
         >>> dim, end_x, end_y = 128, 8, 8
-        >>> freqs_cis_x, freqs_cis_y = compute_axial_cis(dim, end_x, end_y)
-        >>> freqs_cis_x.shape
-        torch.Size([64, 32])
-        >>> freqs_cis_y.shape
-        torch.Size([64, 32])
+        >>> freqs_cis = compute_axial_cis(dim, end_x, end_y)
+        >>> freqs_cis.shape
+        torch.Size([64, 64])
     """
     freqs_x = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
     freqs_y = 1.0 / (theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim))
@@ -156,11 +146,10 @@ def compute_axial_cis(dim: int, end_x: int, end_y: int, theta: float = 10000.0):
 
 
 def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
-    """
-    Reshape frequency tensor for broadcasting with input tensor.
+    """Reshape frequency tensor for broadcasting with input tensor.
 
-    Reshapes a frequency tensor to ensure dimensional compatibility for broadcasting with an input tensor.
-    This function is typically used in positional encoding operations.
+    Reshapes a frequency tensor to ensure dimensional compatibility for broadcasting with an input tensor. This function
+    is typically used in positional encoding operations.
 
     Args:
         freqs_cis (torch.Tensor): Frequency tensor with shape matching the last two dimensions of x.
@@ -185,8 +174,7 @@ def apply_rotary_enc(
     freqs_cis: torch.Tensor,
     repeat_freqs_k: bool = False,
 ):
-    """
-    Apply rotary positional encoding to query and key tensors.
+    """Apply rotary positional encoding to query and key tensors.
 
     This function applies rotary positional encoding (RoPE) to query and key tensors using complex-valued frequency
     components. RoPE is a technique that injects relative position information into self-attention mechanisms.
@@ -194,10 +182,10 @@ def apply_rotary_enc(
     Args:
         xq (torch.Tensor): Query tensor to encode with positional information.
         xk (torch.Tensor): Key tensor to encode with positional information.
-        freqs_cis (torch.Tensor): Complex-valued frequency components for rotary encoding with shape matching the
-            last two dimensions of xq.
-        repeat_freqs_k (bool, optional): Whether to repeat frequency components along sequence length dimension
-            to match key sequence length.
+        freqs_cis (torch.Tensor): Complex-valued frequency components for rotary encoding with shape matching the last
+            two dimensions of xq.
+        repeat_freqs_k (bool, optional): Whether to repeat frequency components along sequence length dimension to match
+            key sequence length.
 
     Returns:
         xq_out (torch.Tensor): Query tensor with rotary positional encoding applied.
@@ -225,9 +213,8 @@ def apply_rotary_enc(
     return xq_out.type_as(xq).to(xq.device), xk_out.type_as(xk).to(xk.device)
 
 
-def window_partition(x, window_size):
-    """
-    Partition input tensor into non-overlapping windows with padding if needed.
+def window_partition(x: torch.Tensor, window_size: int):
+    """Partition input tensor into non-overlapping windows with padding if needed.
 
     Args:
         x (torch.Tensor): Input tensor with shape (B, H, W, C).
@@ -235,7 +222,7 @@ def window_partition(x, window_size):
 
     Returns:
         windows (torch.Tensor): Partitioned windows with shape (B * num_windows, window_size, window_size, C).
-        padded_h_w (Tuple[int, int]): Padded height and width before partition.
+        padded_h_w (tuple[int, int]): Padded height and width before partition.
 
     Examples:
         >>> x = torch.randn(1, 16, 16, 3)
@@ -256,24 +243,23 @@ def window_partition(x, window_size):
     return windows, (Hp, Wp)
 
 
-def window_unpartition(windows, window_size, pad_hw, hw):
-    """
-    Unpartition windowed sequences into original sequences and remove padding.
+def window_unpartition(windows: torch.Tensor, window_size: int, pad_hw: tuple[int, int], hw: tuple[int, int]):
+    """Unpartition windowed sequences into original sequences and remove padding.
 
-    This function reverses the windowing process, reconstructing the original input from windowed segments
-    and removing any padding that was added during the windowing process.
+    This function reverses the windowing process, reconstructing the original input from windowed segments and removing
+    any padding that was added during the windowing process.
 
     Args:
         windows (torch.Tensor): Input tensor of windowed sequences with shape (B * num_windows, window_size,
-            window_size, C), where B is the batch size, num_windows is the number of windows, window_size is
-            the size of each window, and C is the number of channels.
+            window_size, C), where B is the batch size, num_windows is the number of windows, window_size is the size of
+            each window, and C is the number of channels.
         window_size (int): Size of each window.
-        pad_hw (Tuple[int, int]): Padded height and width (Hp, Wp) of the input before windowing.
-        hw (Tuple[int, int]): Original height and width (H, W) of the input before padding and windowing.
+        pad_hw (tuple[int, int]): Padded height and width (Hp, Wp) of the input before windowing.
+        hw (tuple[int, int]): Original height and width (H, W) of the input before padding and windowing.
 
     Returns:
-        (torch.Tensor): Unpartitioned sequences with shape (B, H, W, C), where B is the batch size, H and W
-            are the original height and width, and C is the number of channels.
+        (torch.Tensor): Unpartitioned sequences with shape (B, H, W, C), where B is the batch size, H and W are the
+            original height and width, and C is the number of channels.
 
     Examples:
         >>> windows = torch.rand(32, 8, 8, 64)  # 32 windows of size 8x8 with 64 channels
@@ -295,18 +281,16 @@ def window_unpartition(windows, window_size, pad_hw, hw):
 
 
 def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
-    """
-    Extract relative positional embeddings based on query and key sizes.
+    """Extract relative positional embeddings based on query and key sizes.
 
     Args:
         q_size (int): Size of the query.
         k_size (int): Size of the key.
-        rel_pos (torch.Tensor): Relative position embeddings with shape (L, C), where L is the maximum relative
-            distance and C is the embedding dimension.
+        rel_pos (torch.Tensor): Relative position embeddings with shape (L, C), where L is the maximum relative distance
+            and C is the embedding dimension.
 
     Returns:
-        (torch.Tensor): Extracted positional embeddings according to relative positions, with shape (q_size,
-            k_size, C).
+        (torch.Tensor): Extracted positional embeddings according to relative positions, with shape (q_size, k_size, C).
 
     Examples:
         >>> q_size, k_size = 8, 16
@@ -341,11 +325,10 @@ def add_decomposed_rel_pos(
     q: torch.Tensor,
     rel_pos_h: torch.Tensor,
     rel_pos_w: torch.Tensor,
-    q_size: Tuple[int, int],
-    k_size: Tuple[int, int],
+    q_size: tuple[int, int],
+    k_size: tuple[int, int],
 ) -> torch.Tensor:
-    """
-    Add decomposed Relative Positional Embeddings to the attention map.
+    """Add decomposed Relative Positional Embeddings to the attention map.
 
     This function calculates and applies decomposed Relative Positional Embeddings as described in the MVITv2
     paper. It enhances the attention mechanism by incorporating spatial relationships between query and key
@@ -356,12 +339,12 @@ def add_decomposed_rel_pos(
         q (torch.Tensor): Query tensor in the attention layer with shape (B, q_h * q_w, C).
         rel_pos_h (torch.Tensor): Relative position embeddings for height axis with shape (Lh, C).
         rel_pos_w (torch.Tensor): Relative position embeddings for width axis with shape (Lw, C).
-        q_size (Tuple[int, int]): Spatial sequence size of query q as (q_h, q_w).
-        k_size (Tuple[int, int]): Spatial sequence size of key k as (k_h, k_w).
+        q_size (tuple[int, int]): Spatial sequence size of query q as (q_h, q_w).
+        k_size (tuple[int, int]): Spatial sequence size of key k as (k_h, k_w).
 
     Returns:
-        (torch.Tensor): Updated attention map with added relative positional embeddings, shape
-            (B, q_h * q_w, k_h * k_w).
+        (torch.Tensor): Updated attention map with added relative positional embeddings, shape (B, q_h * q_w, k_h *
+            k_w).
 
     Examples:
         >>> B, C, q_h, q_w, k_h, k_w = 1, 64, 8, 8, 8, 8