rtx-backend-metal: fix swap_dims returning a corrupt strided view

swap_dims copied the buffer but wrote elements back at their ORIGINAL
positions (new_idx was computed with the swapped strides), returning a
stride-swapped non-contiguous tensor. Every other op in this backend —
elementwise kernels, MPS matmul, to_vec — reads raw buffers and ignores
strides, so any transpose consumer (notably the autograd matmul backward,
grad_a = grad_c @ b^T) silently computed on untransposed data. Found via
CPU-vs-Metal gradient parity on the DigiGraph HetGAT: forward matched,
gradients were ~2x off.

swap_dims now physically permutes into a contiguous result (reading
through the input's strides + offset), and reshape asserts contiguity
instead of silently reinterpreting a non-contiguous buffer. 3 new parity
tests incl. matmul-after-transpose (25 total pass on-device).

Co-Authored-By: Claude Fable 5 <[email protected]>

crates/core/rtx-backend-metal/src/ops/shape.rs

@@ -1,4 +1,9 @@
 //! Shape manipulation operations.
+//!
+//! All ops here return **contiguous** tensors. The rest of the backend
+//! (elementwise kernels, MPS matmul, `to_vec`) reads raw buffers and ignores
+//! strides, so a stride-swapped "view" would silently corrupt every
+//! downstream op — `swap_dims` therefore physically permutes the data.
 
 use crate::MetalTensorPrimitive;
 use rtx_metal::MetalBuffer;
@@ -12,6 +17,10 @@ pub fn reshape<const D1: usize, const D2: usize>(
     let old_numel: usize = tensor.shape.iter().product();
     let new_numel: usize = shape.iter().product();
     assert_eq!(old_numel, new_numel, "Total elements must remain the same");
+    assert!(
+        tensor.is_contiguous(),
+        "reshape requires a contiguous tensor (raw-buffer reinterpretation)"
+    );
 
     let strides = MetalTensorPrimitive::<D2>::compute_strides(&shape);
 
@@ -24,7 +33,7 @@ pub fn reshape<const D1: usize, const D2: usize>(
     }
 }
 
-/// Transpose the last two dimensions.
+/// Transpose the last two dimensions (physical, contiguous result).
 pub fn transpose<const D: usize>(tensor: &MetalTensorPrimitive<D>) -> MetalTensorPrimitive<D> {
     if D < 2 {
         return tensor.clone();
@@ -32,7 +41,8 @@ pub fn transpose<const D: usize>(tensor: &MetalTensorPrimitive<D>) -> MetalTenso
     swap_dims(tensor, D - 2, D - 1)
 }
 
-/// Swap two dimensions.
+/// Swap two dimensions, physically permuting the data so the result is
+/// contiguous in the new shape.
 pub fn swap_dims<const D: usize>(
     tensor: &MetalTensorPrimitive<D>,
     dim1: usize,
@@ -45,30 +55,34 @@ pub fn swap_dims<const D: usize>(
     }
 
     let mut new_shape = tensor.shape;
-    let mut new_strides = tensor.strides;
     new_shape.swap(dim1, dim2);
-    new_strides.swap(dim1, dim2);
+    let out_strides = MetalTensorPrimitive::<D>::compute_strides(&new_shape);
 
-    let data = tensor.to_vec();
+    let src = tensor.to_vec();
     let numel = tensor.numel();
     let mut result = vec![0.0f32; numel];
 
-    for i in 0..numel {
+    for (out_pos, r) in result.iter_mut().enumerate() {
-        let mut indices = [0usize; D];
+        // Decompose the contiguous output position into a new_shape
-        let mut remaining = i;
+        // multi-index, map it back (swap) to an input multi-index, and read
+        // through the input's own strides + offset.
+        let mut rem = out_pos;
+        let mut in_off = tensor.offset;
         for d in 0..D {
-            indices[d] = remaining / tensor.strides[d];
+            let id = rem / out_strides[d];
-            remaining %= tensor.strides[d];
+            rem %= out_strides[d];
+            // output dim d corresponds to input dim d, with dim1/dim2 swapped;
+            // the index value along src_dim equals the output index along d
+            let src_dim = if d == dim1 {
+                dim2
+            } else if d == dim2 {
+                dim1
+            } else {
+                d
+            };
+            in_off += id * tensor.strides[src_dim];
         }
-
+        *r = src[in_off];
-        indices.swap(dim1, dim2);
-
-        let mut new_idx = 0;
-        for d in 0..D {
-            new_idx += indices[d] * new_strides[d];
-        }
-
-        result[new_idx] = data[i];
     }
 
     let metal_data = MetalBuffer::from_slice(tensor.device.metal_device(), &result)
@@ -77,7 +91,7 @@ pub fn swap_dims<const D: usize>(
     MetalTensorPrimitive {
         data: Arc::new(metal_data),
         shape: new_shape,
-        strides: new_strides,
+        strides: out_strides,
         device: tensor.device.clone(),
         offset: 0,
     }

crates/core/rtx-backend-metal/tests/backend_parity_tests.rs

@@ -5,7 +5,7 @@
 
 use rtx_backend_metal::MetalDeviceWrapper;
 use rtx_backend_metal::ops::{
-    activation, basic, creation, device as dev_ops, gemm, reduction, unary,
+    activation, basic, creation, device as dev_ops, gemm, reduction, shape, unary,
 };
 
 /// Tolerance for floating-point comparisons.
@@ -608,3 +608,69 @@ fn test_parity_matmul_larger() {
         }
     }
 }
+
+#[test]
+fn test_parity_transpose_2d() {
+    let device = MetalDeviceWrapper::new().expect("Metal device");
+    let data: Vec<f32> = (0..12).map(|x| x as f32).collect(); // [3, 4]
+    let t = creation::from_data(&data, [3, 4], &device);
+    let tt = shape::transpose(&t);
+    // physical transpose: result must be contiguous row-major [4, 3]
+    assert!(tt.is_contiguous(), "transpose must return a contiguous tensor");
+    let got = dev_ops::copy_to_host(&tt);
+    let mut want = vec![0.0f32; 12];
+    for i in 0..3 {
+        for j in 0..4 {
+            want[j * 3 + i] = data[i * 4 + j];
+        }
+    }
+    assert_eq!(got, want);
+    // and transposing back must round-trip
+    let back = dev_ops::copy_to_host(&shape::transpose(&tt));
+    assert_eq!(back, data);
+}
+
+#[test]
+fn test_parity_swap_dims_3d() {
+    let device = MetalDeviceWrapper::new().expect("Metal device");
+    let data: Vec<f32> = (0..24).map(|x| x as f32).collect(); // [2, 3, 4]
+    let t = creation::from_data(&data, [2, 3, 4], &device);
+    let s = shape::swap_dims(&t, 0, 2); // -> [4, 3, 2]
+    assert!(s.is_contiguous());
+    let got = dev_ops::copy_to_host(&s);
+    let mut want = vec![0.0f32; 24];
+    for a in 0..2 {
+        for b in 0..3 {
+            for c in 0..4 {
+                want[c * 6 + b * 2 + a] = data[a * 12 + b * 4 + c];
+            }
+        }
+    }
+    assert_eq!(got, want);
+}
+
+#[test]
+fn test_parity_matmul_after_transpose() {
+    // the exact pattern the autograd matmul backward uses:
+    // grad_a = grad_c @ b^T — a transposed operand must feed MPS correctly.
+    let device = MetalDeviceWrapper::new().expect("Metal device");
+    let a: Vec<f32> = (0..6).map(|x| x as f32 + 1.0).collect(); // [2, 3]
+    let b: Vec<f32> = (0..12).map(|x| (x as f32) * 0.5 - 2.0).collect(); // [4, 3] -> b^T [3, 4]
+    let ta = creation::from_data(&a, [2, 3], &device);
+    let tb = creation::from_data(&b, [4, 3], &device);
+    let c = gemm::matmul(&ta, &shape::transpose(&tb));
+    let got = dev_ops::copy_to_host(&c);
+    let mut want = vec![0.0f32; 8];
+    for i in 0..2 {
+        for j in 0..4 {
+            let mut s = 0.0f32;
+            for k in 0..3 {
+                s += a[i * 3 + k] * b[j * 3 + k];
+            }
+            want[i * 4 + j] = s;
+        }
+    }
+    for (g, w) in got.iter().zip(want.iter()) {
+        assert!((g - w).abs() < 1e-4, "matmul-after-transpose mismatch: {got:?} vs {want:?}");
+    }
+}