diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index ceb39ff..54f1e1c 100644 --- a/crates/mlxcore/src/array.rs +++ b/crates/mlxcore/src/array.rs @@ -181,6 +181,46 @@ impl Array { self.unary_op(stream, sys::mlx_negative) } + /// Elementwise natural logarithm. + pub fn log(&self, stream: &Stream) -> Result { + self.unary_op(stream, sys::mlx_log) + } + + /// Elementwise sine. + pub fn sin(&self, stream: &Stream) -> Result { + self.unary_op(stream, sys::mlx_sin) + } + + /// Elementwise cosine. + pub fn cos(&self, stream: &Stream) -> Result { + self.unary_op(stream, sys::mlx_cos) + } + + /// Elementwise square. + pub fn square(&self, stream: &Stream) -> Result { + self.unary_op(stream, sys::mlx_square) + } + + /// Elementwise hyperbolic tangent. + pub fn tanh(&self, stream: &Stream) -> Result { + self.unary_op(stream, sys::mlx_tanh) + } + + /// Elementwise power: `self ** other`. + pub fn power(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_power) + } + + /// Elementwise maximum of two arrays. + pub fn maximum(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_maximum) + } + + /// Elementwise minimum of two arrays. + pub fn minimum(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_minimum) + } + /// Sum of all elements, returning a scalar array. /// /// With `keepdims == false` the result is 0-dimensional. @@ -195,6 +235,46 @@ impl Array { self.reduce_op(keepdims, stream, sys::mlx_mean) } + /// Sum over the given axes. + /// + /// With `keepdims == false` the reduced axes are removed; otherwise they + /// are kept with size 1. + pub fn sum_axes(&self, axes: &[i32], keepdims: bool, stream: &Stream) -> Result { + self.reduce_axes_op(axes, keepdims, stream, sys::mlx_sum_axes) + } + + /// Mean over the given axes. + /// + /// With `keepdims == false` the reduced axes are removed; otherwise they + /// are kept with size 1. + pub fn mean_axes(&self, axes: &[i32], keepdims: bool, stream: &Stream) -> Result { + self.reduce_axes_op(axes, keepdims, stream, sys::mlx_mean_axes) + } + + /// Maximum over the given axes. + /// + /// With `keepdims == false` the reduced axes are removed; otherwise they + /// are kept with size 1. + pub fn max_axes(&self, axes: &[i32], keepdims: bool, stream: &Stream) -> Result { + self.reduce_axes_op(axes, keepdims, stream, sys::mlx_max_axes) + } + + /// Minimum over the given axes. + /// + /// With `keepdims == false` the reduced axes are removed; otherwise they + /// are kept with size 1. + pub fn min_axes(&self, axes: &[i32], keepdims: bool, stream: &Stream) -> Result { + self.reduce_axes_op(axes, keepdims, stream, sys::mlx_min_axes) + } + + /// Product over the given axes. + /// + /// With `keepdims == false` the reduced axes are removed; otherwise they + /// are kept with size 1. + pub fn prod_axes(&self, axes: &[i32], keepdims: bool, stream: &Stream) -> Result { + self.reduce_axes_op(axes, keepdims, stream, sys::mlx_prod_axes) + } + /// Shared plumbing for `res = op(a, b, stream)` binary ops. fn binary_op( &self, @@ -241,6 +321,48 @@ impl Array { Self::from_op(out, status) } + /// Shared plumbing for `res = op(a, axes, axes_num, keepdims, stream)` + /// reductions over specific axes. + fn reduce_axes_op( + &self, + axes: &[i32], + keepdims: bool, + stream: &Stream, + op: unsafe extern "C" fn( + *mut sys::mlx_array, + sys::mlx_array, + *const i32, + usize, + bool, + sys::mlx_stream, + ) -> i32, + ) -> Result { + error::install(); + // For an empty slice `as_ptr()` is non-null but dangling; pass an + // explicit null pointer so the FFI call never hands C a bogus pointer + // even if it were to dereference it with `axes_num == 0`. + let axes_ptr = if axes.is_empty() { + std::ptr::null() + } else { + axes.as_ptr() + }; + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: `axes_ptr`/`axes.len()` describe a valid slice (or null/0) for + // the duration of the call; all handles are valid; `op` writes the + // result into `out`. + let status = unsafe { + op( + &mut out, + self.handle, + axes_ptr, + axes.len(), + keepdims, + stream.as_raw(), + ) + }; + Self::from_op(out, status) + } + /// Wraps an op's `out` handle and status code into a `Result`. /// /// On failure, frees the (unused) `out` handle and returns the captured @@ -397,6 +519,48 @@ mod tests { ); } + #[test] + fn more_unary_ops() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3]); + assert_eq!(a.square(&s).unwrap().to_vec::(), vec![1.0, 4.0, 9.0]); + + // log(1) == 0, and log(e) ~= 1. + let b = Array::from_slice(&[1.0f32, std::f32::consts::E], &[2]); + let logged = b.log(&s).unwrap().to_vec::(); + assert!((logged[0]).abs() < 1e-6); + assert!((logged[1] - 1.0).abs() < 1e-6); + + // sin(0) == 0, cos(0) == 1, tanh(0) == 0. + let z = Array::from_slice(&[0.0f32], &[1]); + assert!(z.sin(&s).unwrap().item::().abs() < 1e-6); + assert!((z.cos(&s).unwrap().item::() - 1.0).abs() < 1e-6); + assert!(z.tanh(&s).unwrap().item::().abs() < 1e-6); + } + + #[test] + fn more_binary_ops() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3]); + let b = Array::from_slice(&[3.0f32, 2.0, 1.0], &[3]); + + assert_eq!( + a.maximum(&b, &s).unwrap().to_vec::(), + vec![3.0, 2.0, 3.0] + ); + assert_eq!( + a.minimum(&b, &s).unwrap().to_vec::(), + vec![1.0, 2.0, 1.0] + ); + + let base = Array::from_slice(&[2.0f32, 3.0], &[2]); + let exp = Array::from_slice(&[3.0f32, 2.0], &[2]); + assert_eq!( + base.power(&exp, &s).unwrap().to_vec::(), + vec![8.0, 9.0] + ); + } + #[test] fn reductions_produce_scalars() { let s = Stream::cpu(); @@ -418,6 +582,56 @@ mod tests { assert_eq!(sum.item::(), 10.0); } + #[test] + fn axis_reductions() { + let s = Stream::cpu(); + // [[1, 2, 3], + // [4, 5, 6]] + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[2, 3]); + + // Sum over axis 0 (rows) -> [5, 7, 9], shape [3]. + let col_sums = a.sum_axes(&[0], false, &s).unwrap(); + assert_eq!(col_sums.shape(), vec![3]); + assert_eq!(col_sums.to_vec::(), vec![5.0, 7.0, 9.0]); + + // Sum over axis 1 (cols) -> [6, 15], shape [2]. + let row_sums = a.sum_axes(&[1], false, &s).unwrap(); + assert_eq!(row_sums.to_vec::(), vec![6.0, 15.0]); + + // keepdims keeps the reduced axis as size 1. + let kept = a.sum_axes(&[1], true, &s).unwrap(); + assert_eq!(kept.shape(), vec![2, 1]); + + // max / min over axis 0; mean / prod over axis 1. + assert_eq!( + a.max_axes(&[0], false, &s).unwrap().to_vec::(), + vec![4.0, 5.0, 6.0] + ); + assert_eq!( + a.min_axes(&[0], false, &s).unwrap().to_vec::(), + vec![1.0, 2.0, 3.0] + ); + assert_eq!( + a.mean_axes(&[1], false, &s).unwrap().to_vec::(), + vec![2.0, 5.0] + ); + assert_eq!( + a.prod_axes(&[1], false, &s).unwrap().to_vec::(), + vec![6.0, 120.0] + ); + } + + #[test] + fn empty_axes_reduction_is_noop() { + // Reducing over no axes must not pass a dangling pointer to C; MLX + // treats it as an identity that leaves the values unchanged. + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]); + let r = a.sum_axes(&[], false, &s).unwrap(); + assert_eq!(r.shape(), vec![2, 2]); + assert_eq!(r.to_vec::(), vec![1.0, 2.0, 3.0, 4.0]); + } + #[test] fn incompatible_shapes_return_err() { let s = Stream::cpu();