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 {
-        let mut indices = [0usize; D];
-        let mut remaining = i;
+    for (out_pos, r) in result.iter_mut().enumerate() {
+        // Decompose the contiguous output position into a new_shape
+        // 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];
-            remaining %= tensor.strides[d];
+            let id = rem / out_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];
         }
-
-        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];
+        *r = src[in_off];
     }
 
     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,
     }