diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index 54f1e1c..9a9cc23 100644 --- a/crates/mlxcore/src/array.rs +++ b/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::()`. 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(&self) -> Vec { - 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 { + 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 { 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 { + self.shape_op(shape, stream, sys::mlx_reshape) + } + + /// Broadcasts the array to `shape`. + pub fn broadcast_to(&self, shape: &[i32], stream: &Stream) -> Result { + 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 { + self.unary_op(stream, sys::mlx_transpose) + } + + /// Removes all axes of length 1. + pub fn squeeze(&self, stream: &Stream) -> Result { + 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 { + 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 { + 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::(), 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::(), 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::(), 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::(), 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();