No VJP for mask or sinks in attention (#2909)

2025-12-16 01:49:05 +08:00 · 2025-12-13 19:48:39 -08:00
parent 47d2505ea9
commit 1b591ec736
1 changed files with 5 additions and 0 deletions
--- a/mlx/fast.cpp
+++ b/mlx/fast.cpp
@@ -880,6 +880,11 @@ std::vector<array> ScaledDotProductAttention::vjp(

  std::vector<array> returned_vjps;
  for (int arg : argnums) {
+    if (arg >= 3) {
+      throw std::invalid_argument(
+          "[scale_dot_product_attention] Does not support VJP with respect "
+          " to mask or attention sinks.");
+    }
    returned_vjps.push_back(std::move(vjps[arg]));
  }
  return returned_vjps;