feat: add shape manipulation ops on Array

crates/mlxcore/src/array.rs

@@ -115,32 +115,60 @@ impl Array {
     /// The element type `T` selects the accessor at compile time, e.g.
     /// `a.to_vec::<f32>()`. Evaluates the array first.
     ///
+    /// Non-contiguous arrays (e.g. from [`transpose`](Self::transpose) or
+    /// [`broadcast_to`](Self::broadcast_to)) are first materialized into a
+    /// row-contiguous copy, so the result always reflects the logical
+    /// (row-major) element order rather than the raw storage buffer.
+    ///
     /// # Panics
-    /// Panics if `T::DTYPE` does not match the array's dtype.
+    /// Panics if `T::DTYPE` does not match the array's dtype, or if
+    /// making the array contiguous fails.
     pub fn to_vec<T: ArrayElement>(&self) -> Vec<T> {
-        self.eval();
+        // Materialize a row-contiguous copy: strided views (transpose) and
+        // stride-0 views (broadcast) don't lay their logical elements out
+        // contiguously in the storage buffer, so reading the raw pointer would
+        // return storage order (or read past the real data). `mlx_contiguous`
+        // produces a dense buffer whose memory order matches the logical order.
+        //
+        // Run it on the CPU stream: this is a host-side data-marshalling step
+        // (we're about to read the buffer from Rust), and it keeps `to_vec` off
+        // the GPU stream so concurrent callers don't contend on Metal.
+        let contiguous = self.contiguous(&Stream::cpu()).unwrap_or_else(|e| {
+            panic!("to_vec: failed to make array contiguous: {e}");
+        });
+        contiguous.eval();
+
         // SAFETY: mlx_array_dtype reads a valid handle.
-        let dtype = unsafe { sys::mlx_array_dtype(self.handle) };
+        let dtype = unsafe { sys::mlx_array_dtype(contiguous.handle) };
         assert_eq!(
             dtype,
             T::DTYPE,
             "array dtype does not match requested element type"
         );
-        let len = self.size();
+        let len = contiguous.size();
         // `from_raw_parts` requires a non-null, aligned pointer even for a
         // zero-length slice, but mlx may return null for an empty array.
         if len == 0 {
             return Vec::new();
         }
-        // SAFETY: dtype matches `T` (checked above), so mlx guarantees `len`
-        // contiguous, aligned `T` at `ptr`, valid until the array is mutated or
-        // freed. We only read (and copy out of) the slice within this call, so
-        // the borrow cannot outlive the buffer. `T: Copy`, so `to_vec` is a
-        // single bulk copy rather than `len` individual derefs.
-        let ptr = unsafe { T::data_ptr(self.handle) };
+        // SAFETY: dtype matches `T` (checked above) and `contiguous` is dense,
+        // so mlx guarantees `len` contiguous, aligned `T` at `ptr`, valid until
+        // `contiguous` is dropped at the end of this function. We copy out of
+        // the slice before that. `T: Copy`, so this is a single bulk copy.
+        let ptr = unsafe { T::data_ptr(contiguous.handle) };
         unsafe { std::slice::from_raw_parts(ptr, len) }.to_vec()
     }
 
+    /// Returns a row-contiguous copy (or the same array if already dense).
+    pub fn contiguous(&self, stream: &Stream) -> Result<Array> {
+        error::install();
+        let mut out = unsafe { sys::mlx_array_new() };
+        // SAFETY: handle/stream valid; `allow_col_major = false` forces
+        // row-major; result written into `out`.
+        let status = unsafe { sys::mlx_contiguous(&mut out, self.handle, false, stream.as_raw()) };
+        Self::from_op(out, status)
+    }
+
     /// Elementwise addition: `self + other`.
     pub fn add(&self, other: &Array, stream: &Stream) -> Result<Array> {
         self.binary_op(other, stream, sys::mlx_add)
@@ -275,6 +303,73 @@ impl Array {
         self.reduce_axes_op(axes, keepdims, stream, sys::mlx_prod_axes)
     }
 
+    /// Returns a new array with the same data reinterpreted as `shape`.
+    ///
+    /// The product of `shape` must equal [`size`](Self::size).
+    pub fn reshape(&self, shape: &[i32], stream: &Stream) -> Result<Array> {
+        self.shape_op(shape, stream, sys::mlx_reshape)
+    }
+
+    /// Broadcasts the array to `shape`.
+    pub fn broadcast_to(&self, shape: &[i32], stream: &Stream) -> Result<Array> {
+        self.shape_op(shape, stream, sys::mlx_broadcast_to)
+    }
+
+    /// Reverses the order of all axes (a full transpose).
+    pub fn transpose(&self, stream: &Stream) -> Result<Array> {
+        self.unary_op(stream, sys::mlx_transpose)
+    }
+
+    /// Removes all axes of length 1.
+    pub fn squeeze(&self, stream: &Stream) -> Result<Array> {
+        self.unary_op(stream, sys::mlx_squeeze)
+    }
+
+    /// Inserts a new axis of length 1 at position `axis`.
+    pub fn expand_dims(&self, axis: i32, stream: &Stream) -> Result<Array> {
+        error::install();
+        let mut out = unsafe { sys::mlx_array_new() };
+        // SAFETY: handle/stream are valid; `op` writes the result into `out`.
+        let status = unsafe { sys::mlx_expand_dims(&mut out, self.handle, axis, stream.as_raw()) };
+        Self::from_op(out, status)
+    }
+
+    /// Shared plumbing for `res = op(a, shape, shape_num, stream)` shape ops.
+    fn shape_op(
+        &self,
+        shape: &[i32],
+        stream: &Stream,
+        op: unsafe extern "C" fn(
+            *mut sys::mlx_array,
+            sys::mlx_array,
+            *const i32,
+            usize,
+            sys::mlx_stream,
+        ) -> i32,
+    ) -> Result<Array> {
+        error::install();
+        // For an empty slice `as_ptr()` is non-null but dangling; pass an
+        // explicit null pointer so C never receives a bogus pointer.
+        let shape_ptr = if shape.is_empty() {
+            std::ptr::null()
+        } else {
+            shape.as_ptr()
+        };
+        let mut out = unsafe { sys::mlx_array_new() };
+        // SAFETY: `shape_ptr`/`shape.len()` describe a valid slice (or null/0)
+        // for the call; all handles are valid; `op` writes into `out`.
+        let status = unsafe {
+            op(
+                &mut out,
+                self.handle,
+                shape_ptr,
+                shape.len(),
+                stream.as_raw(),
+            )
+        };
+        Self::from_op(out, status)
+    }
+
     /// Shared plumbing for `res = op(a, b, stream)` binary ops.
     fn binary_op(
         &self,
@@ -632,6 +727,46 @@ mod tests {
         assert_eq!(r.to_vec::<f32>(), vec![1.0, 2.0, 3.0, 4.0]);
     }
 
+    #[test]
+    fn reshape_changes_shape_not_data() {
+        let s = Stream::cpu();
+        let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
+        let r = a.reshape(&[3, 2], &s).unwrap();
+        assert_eq!(r.shape(), vec![3, 2]);
+        assert_eq!(r.to_vec::<f32>(), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
+    }
+
+    #[test]
+    fn transpose_reverses_axes() {
+        let s = Stream::cpu();
+        // [[1, 2, 3],
+        //  [4, 5, 6]]  ->  [[1, 4], [2, 5], [3, 6]]
+        let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]);
+        let t = a.transpose(&s).unwrap();
+        assert_eq!(t.shape(), vec![3, 2]);
+        assert_eq!(t.to_vec::<f32>(), vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
+    }
+
+    #[test]
+    fn broadcast_to_expands() {
+        let s = Stream::cpu();
+        let a = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3]);
+        let b = a.broadcast_to(&[2, 3], &s).unwrap();
+        assert_eq!(b.shape(), vec![2, 3]);
+        assert_eq!(b.to_vec::<f32>(), vec![1.0, 2.0, 3.0, 1.0, 2.0, 3.0]);
+    }
+
+    #[test]
+    fn squeeze_and_expand_dims() {
+        let s = Stream::cpu();
+        let a = Array::from_slice(&[1.0f32, 2.0, 3.0], &[1, 3, 1]);
+        let sq = a.squeeze(&s).unwrap();
+        assert_eq!(sq.shape(), vec![3]);
+
+        let ex = sq.expand_dims(0, &s).unwrap();
+        assert_eq!(ex.shape(), vec![1, 3]);
+    }
+
     #[test]
     fn incompatible_shapes_return_err() {
         let s = Stream::cpu();