fix(vision): capture the functional slice_update return in DeepStack and MRoPE-section loops

src/models/paddleocr_vl.rs

@@ -188,7 +188,15 @@ impl MRoPE {
                 &[src_dim + 1, batch, seq_len, end],
             );
             let section = mlxcel_core::squeeze_axis(&section, 0);
-            mlxcel_core::slice_update(&result, &section, &[0, 0, offset], &[batch, seq_len, end]);
+            // `slice_update` is functional; capture the return or the H/W
+            // sections are silently dropped and every dim keeps the temporal
+            // frequencies (issue #650).
+            result = mlxcel_core::slice_update(
+                &result,
+                &section,
+                &[0, 0, offset],
+                &[batch, seq_len, end],
+            );
             offset = end;
         }
         result
@@ -704,3 +712,34 @@ impl mlxcel_core::generate::LanguageModel for PaddleOcrTextModel {
         self.eos_token_ids.clone()
     }
 }
+
+#[cfg(test)]
+mod mrope_section_tests {
+    use super::*;
+
+    /// Regression for issue #650: `apply_mrope` must actually write the H and
+    /// W sections (the functional `slice_update` return was dropped, leaving
+    /// every dim on the temporal frequencies).
+    #[test]
+    fn apply_mrope_writes_h_and_w_sections() {
+        let rope = MRoPE::new(12, 10_000.0, vec![2, 2, 2]);
+        // freqs [3, batch=1, seq=1, half=6]: axis T = 1.0, H = 2.0, W = 3.0.
+        let mut data = Vec::new();
+        for axis_val in [1.0f32, 2.0, 3.0] {
+            data.extend(std::iter::repeat_n(axis_val, 6));
+        }
+        let freqs = mlxcel_core::from_slice_f32(&data, &[3, 1, 1, 6]);
+        let mixed = rope.apply_mrope(&freqs);
+        mlxcel_core::eval(&mixed);
+        let expected = [1.0f32, 1.0, 2.0, 2.0, 3.0, 3.0];
+        for (i, &want) in expected.iter().enumerate() {
+            let v = mlxcel_core::slice(&mixed, &[0, 0, i as i32], &[1, 1, i as i32 + 1]);
+            mlxcel_core::eval(&v);
+            let got = mlxcel_core::item_f32(&v);
+            assert!(
+                (got - want).abs() < 1e-6,
+                "dim {i}: expected axis value {want}, got {got} (H/W section write dropped?)"
+            );
+        }
+    }
+}

src/models/qwen2_vl.rs

@@ -202,7 +202,15 @@ impl MRoPE {
             let section = mlxcel_core::squeeze_axis(&section, 0);
 
             // Replace in result via slice_update
-            mlxcel_core::slice_update(&result, &section, &[0, 0, offset], &[batch, seq_len, end]);
+            // `slice_update` is functional; capture the return or the H/W
+            // sections are silently dropped and every dim keeps the temporal
+            // frequencies (issue #650).
+            result = mlxcel_core::slice_update(
+                &result,
+                &section,
+                &[0, 0, offset],
+                &[batch, seq_len, end],
+            );
 
             offset = end;
         }
@@ -798,3 +806,34 @@ impl mlxcel_core::generate::LanguageModel for Qwen2VLModel {
         vec![151645, 151643] // Qwen2 EOS tokens
     }
 }
+
+#[cfg(test)]
+mod mrope_section_tests {
+    use super::*;
+
+    /// Regression for issue #650: `apply_mrope` must actually write the H and
+    /// W sections (the functional `slice_update` return was dropped, leaving
+    /// every dim on the temporal frequencies).
+    #[test]
+    fn apply_mrope_writes_h_and_w_sections() {
+        let rope = MRoPE::new(12, 10_000.0, vec![2, 2, 2]);
+        // freqs [3, batch=1, seq=1, half=6]: axis T = 1.0, H = 2.0, W = 3.0.
+        let mut data = Vec::new();
+        for axis_val in [1.0f32, 2.0, 3.0] {
+            data.extend(std::iter::repeat_n(axis_val, 6));
+        }
+        let freqs = mlxcel_core::from_slice_f32(&data, &[3, 1, 1, 6]);
+        let mixed = rope.apply_mrope(&freqs);
+        mlxcel_core::eval(&mixed);
+        let expected = [1.0f32, 1.0, 2.0, 2.0, 3.0, 3.0];
+        for (i, &want) in expected.iter().enumerate() {
+            let v = mlxcel_core::slice(&mixed, &[0, 0, i as i32], &[1, 1, i as i32 + 1]);
+            mlxcel_core::eval(&v);
+            let got = mlxcel_core::item_f32(&v);
+            assert!(
+                (got - want).abs() < 1e-6,
+                "dim {i}: expected axis value {want}, got {got} (H/W section write dropped?)"
+            );
+        }
+    }
+}

src/models/qwen3_vl.rs

@@ -788,14 +788,34 @@ impl Qwen3VLModel {
             // Build result by scatter
             // For simplicity, use the full tensor approach:
             // result = h.copy(), then update positions
-            let result = mlxcel_core::copy(&batch_h);
-            for (local_idx, &pos) in image_positions.iter().enumerate() {
-                let val = mlxcel_core::slice(
+            // `slice_update` is functional (returns a new array); the result
+            // MUST be captured or the write is lost (issue #650, same class as
+            // the Granite 4 Vision blindness fixed by `inject_at_positions`).
+            // Contiguous image runs (the single-image case) collapse to one
+            // slice_update; disjoint runs fall back to per-row.
+            let mut result = mlxcel_core::copy(&batch_h);
+            let contiguous = image_positions
+                .iter()
+                .enumerate()
+                .all(|(i, &p)| p == image_positions[0] + i as i32);
+            if contiguous {
+                let start = image_positions[0];
+                result = mlxcel_core::slice_update(
+                    &result,
                     &updated_vals,
-                    &[local_idx as i32, 0],
-                    &[local_idx as i32 + 1, h_shape[2]],
+                    &[start, 0],
+                    &[start + n_img, h_shape[2]],
                 );
-                mlxcel_core::slice_update(&result, &val, &[pos, 0], &[pos + 1, h_shape[2]]);
+            } else {
+                for (local_idx, &pos) in image_positions.iter().enumerate() {
+                    let val = mlxcel_core::slice(
+                        &updated_vals,
+                        &[local_idx as i32, 0],
+                        &[local_idx as i32 + 1, h_shape[2]],
+                    );
+                    result =
+                        mlxcel_core::slice_update(&result, &val, &[pos, 0], &[pos + 1, h_shape[2]]);
+                }
             }
 
             mlxcel_core::expand_dims(&result, 0)
@@ -1042,3 +1062,37 @@ impl mlxcel_core::generate::LanguageModel for Qwen3VLModel {
         vec![151645, 151643] // Qwen EOS tokens
     }
 }
+
+#[cfg(test)]
+mod deepstack_injection_tests {
+    use super::*;
+
+    /// Regression for issue #650: `deepstack_process` must add the visual
+    /// embeds at image positions (the functional `slice_update` return was
+    /// dropped, silently skipping the whole injection).
+    #[test]
+    fn deepstack_process_adds_visual_embeds_at_image_positions() {
+        let hidden = 4i32;
+        let seq = 5i32;
+        let h = mlxcel_core::full_f32(&[1, seq, hidden], 1.0, mlxcel_core::dtype::FLOAT32);
+        // Image positions 1..=3 (contiguous run).
+        let mask_vals = [0.0f32, 1.0, 1.0, 1.0, 0.0];
+        let mask = mlxcel_core::from_slice_f32(&mask_vals, &[1, seq]);
+        let mask = mlxcel_core::astype(&mask, mlxcel_core::dtype::BOOL);
+        let visual = mlxcel_core::full_f32(&[3, hidden], 10.0, mlxcel_core::dtype::FLOAT32);
+
+        let out = Qwen3VLModel::deepstack_process(&h, &mask, &visual);
+        mlxcel_core::eval(&out);
+
+        for pos in 0..seq {
+            let v = mlxcel_core::slice(&out, &[0, pos, 0], &[1, pos + 1, 1]);
+            mlxcel_core::eval(&v);
+            let got = mlxcel_core::item_f32(&v);
+            let want = if (1..=3).contains(&pos) { 11.0 } else { 1.0 };
+            assert!(
+                (got - want).abs() < 1e-6,
+                "pos {pos}: expected {want}, got {got} (injection dropped?)"
+            );
+        }
+    }
+}

src/models/qwen3_vl_moe.rs

@@ -981,14 +981,34 @@ impl Qwen3VLMoeModel {
             let visual_slice = mlxcel_core::slice(visual_embeds, &[0, 0], &[n_img, h_shape[2]]);
             let updated_vals = mlxcel_core::add(&current_vals, &visual_slice);
 
-            let result = mlxcel_core::copy(&batch_h);
-            for (local_idx, &pos) in image_positions.iter().enumerate() {
-                let val = mlxcel_core::slice(
+            // `slice_update` is functional (returns a new array); the result
+            // MUST be captured or the write is lost (issue #650, same class as
+            // the Granite 4 Vision blindness fixed by `inject_at_positions`).
+            // Contiguous image runs (the single-image case) collapse to one
+            // slice_update; disjoint runs fall back to per-row.
+            let mut result = mlxcel_core::copy(&batch_h);
+            let contiguous = image_positions
+                .iter()
+                .enumerate()
+                .all(|(i, &p)| p == image_positions[0] + i as i32);
+            if contiguous {
+                let start = image_positions[0];
+                result = mlxcel_core::slice_update(
+                    &result,
                     &updated_vals,
-                    &[local_idx as i32, 0],
-                    &[local_idx as i32 + 1, h_shape[2]],
+                    &[start, 0],
+                    &[start + n_img, h_shape[2]],
                 );
-                mlxcel_core::slice_update(&result, &val, &[pos, 0], &[pos + 1, h_shape[2]]);
+            } else {
+                for (local_idx, &pos) in image_positions.iter().enumerate() {
+                    let val = mlxcel_core::slice(
+                        &updated_vals,
+                        &[local_idx as i32, 0],
+                        &[local_idx as i32 + 1, h_shape[2]],
+                    );
+                    result =
+                        mlxcel_core::slice_update(&result, &val, &[pos, 0], &[pos + 1, h_shape[2]]);
+                }
             }
 
             mlxcel_core::expand_dims(&result, 0)
@@ -1234,3 +1254,37 @@ impl mlxcel_core::generate::LanguageModel for Qwen3VLMoeModel {
         vec![151645, 151643] // Qwen EOS tokens
     }
 }
+
+#[cfg(test)]
+mod deepstack_injection_tests {
+    use super::*;
+
+    /// Regression for issue #650: `deepstack_process` must add the visual
+    /// embeds at image positions (the functional `slice_update` return was
+    /// dropped, silently skipping the whole injection).
+    #[test]
+    fn deepstack_process_adds_visual_embeds_at_image_positions() {
+        let hidden = 4i32;
+        let seq = 5i32;
+        let h = mlxcel_core::full_f32(&[1, seq, hidden], 1.0, mlxcel_core::dtype::FLOAT32);
+        // Image positions 1..=3 (contiguous run).
+        let mask_vals = [0.0f32, 1.0, 1.0, 1.0, 0.0];
+        let mask = mlxcel_core::from_slice_f32(&mask_vals, &[1, seq]);
+        let mask = mlxcel_core::astype(&mask, mlxcel_core::dtype::BOOL);
+        let visual = mlxcel_core::full_f32(&[3, hidden], 10.0, mlxcel_core::dtype::FLOAT32);
+
+        let out = Qwen3VLMoeModel::deepstack_process(&h, &mask, &visual);
+        mlxcel_core::eval(&out);
+
+        for pos in 0..seq {
+            let v = mlxcel_core::slice(&out, &[0, pos, 0], &[1, pos + 1, 1]);
+            mlxcel_core::eval(&v);
+            let got = mlxcel_core::item_f32(&v);
+            let want = if (1..=3).contains(&pos) { 11.0 } else { 1.0 };
+            assert!(
+                (got - want).abs() < 1e-6,
+                "pos {pos}: expected {want}, got {got} (injection dropped?)"
+            );
+        }
+    }
+}