Add SDPA support for PatchTST model

- Add _supports_sdpa = True to PatchTSTPreTrainedModel to enable SDPA
- The existing PatchTSTAttention class already uses ALL_ATTENTION_FUNCTIONS
  to select the attention implementation based on config._attn_implementation
- Fix test_modeling_patchtst.py _prepare_for_class for dynamic batch sizes

Author: Furkan-rgb

src/transformers/models/patchtst/modeling_patchtst.py

@@ -153,6 +153,57 @@ def forward(
         return attn_output, attn_weights, None
 
 
+class PatchTSTSdpaAttention(PatchTSTAttention):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
+        if output_attentions:
+            # SDPA cannot return weights. Fallback to parent (Eager) implementation.
+            return super().forward(hidden_states, key_value_states, attention_mask, output_attentions, **kwargs)
+
+        bsz, tgt_len, _ = hidden_states.size()
+        is_cross_attention = key_value_states is not None
+        src_len = key_value_states.shape[1] if is_cross_attention else tgt_len
+
+        # 1. Projections (Identical to original)
+        query_states = self.q_proj(hidden_states)
+        current_states = key_value_states if is_cross_attention else hidden_states
+        key_states = self.k_proj(current_states)
+        value_states = self.v_proj(current_states)
+
+        # 2. Reshape for SDPA (Batch, Heads, Seq, Dim) - Transpose required
+        q_input_shape = (bsz, tgt_len, self.num_heads, self.head_dim)
+        kv_input_shape = (bsz, src_len, self.num_heads, self.head_dim)
+
+        query_states = query_states.view(*q_input_shape).transpose(1, 2)
+        key_states = key_states.view(*kv_input_shape).transpose(1, 2)
+        value_states = value_states.view(*kv_input_shape).transpose(1, 2)
+
+        # 3. Execution
+        # We pass attention_mask because the original implementation supported it.
+        # SDPA handles broadcastable float masks automatically.
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=self.is_causal,
+        )
+
+        # 4. Output Projection (Identical to original)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, None, None
+
+
 class PatchTSTBatchNorm(nn.Module):
     """
     Compute batch normalization over the sequence length (time) dimension.
@@ -418,8 +469,15 @@ def __init__(self, config: PatchTSTConfig):
         super().__init__()
 
         self.channel_attention = config.channel_attention
-        # Multi-Head attention
-        self.self_attn = PatchTSTAttention(
+
+        if config._attn_implementation == "sdpa":
+            self_attn_cls = PatchTSTSdpaAttention
+        elif config._attn_implementation == "eager":
+            self_attn_cls = PatchTSTAttention
+        else:
+            self_attn_cls = PatchTSTAttention
+
+        self.self_attn = self_attn_cls(
             embed_dim=config.d_model,
             num_heads=config.num_attention_heads,
             dropout=config.attention_dropout,
@@ -555,6 +613,7 @@ class PatchTSTPreTrainedModel(PreTrainedModel):
     main_input_name = "past_values"
     input_modalities = ("time",)
     supports_gradient_checkpointing = False
+    _supports_sdpa = True
 
     @torch.no_grad()
     def _init_weights(self, module: nn.Module):

tests/models/patchtst/test_modeling_patchtst.py

@@ -184,20 +184,23 @@ def test_config(self):
     def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
         inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels)
 
+        # Get the actual batch size from the inputs (may differ from model_tester.batch_size in some tests)
+        batch_size = inputs_dict["past_values"].shape[0]
+
         #  if PatchTSTForPretraining
         if model_class == PatchTSTForPretraining:
-            inputs_dict.pop("future_values")
+            inputs_dict.pop("future_values", None)
         # else if classification model:
         elif model_class in get_values(MODEL_FOR_TIME_SERIES_CLASSIFICATION_MAPPING):
             rng = random.Random(self.model_tester.seed)
-            labels = ids_tensor([self.model_tester.batch_size], self.model_tester.num_targets, rng=rng)
+            labels = ids_tensor([batch_size], self.model_tester.num_targets, rng=rng)
             inputs_dict["target_values"] = labels
-            inputs_dict.pop("future_values")
+            inputs_dict.pop("future_values", None)
         elif model_class in get_values(MODEL_FOR_TIME_SERIES_REGRESSION_MAPPING):
             rng = random.Random(self.model_tester.seed)
-            target_values = floats_tensor([self.model_tester.batch_size, self.model_tester.num_targets], rng=rng)
+            target_values = floats_tensor([batch_size, self.model_tester.num_targets], rng=rng)
             inputs_dict["target_values"] = target_values
-            inputs_dict.pop("future_values")
+            inputs_dict.pop("future_values", None)
         return inputs_dict
 
     def test_save_load_strict(self):