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

2025-12-13 19:48:39 -08:00
parent 47d2505ea9
commit 1b591ec736
1 changed files with 5 additions and 0 deletions
@@ -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;