From 0892e5acad3458510966a5dd6a54a1f731a05c99 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 15:08:28 +0100 Subject: [PATCH 1/6] support more ops Signed-off-by: kerthcet --- crates/mlxcore/src/array.rs | 81 +++++++++++++++++++++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index ceb39ff..f653419 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. @@ -397,6 +437,47 @@ 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. + 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); + } + + #[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(); From 704e46097a89a69040db12002c857a7e2443f06b Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 16:11:58 +0100 Subject: [PATCH 2/6] add axis-specific reductions Signed-off-by: kerthcet --- crates/mlxcore/src/array.rs | 100 ++++++++++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index f653419..3d03f0a 100644 --- a/crates/mlxcore/src/array.rs +++ b/crates/mlxcore/src/array.rs @@ -235,6 +235,34 @@ 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. + 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. + 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. + 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. + 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, @@ -281,6 +309,39 @@ 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(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: `axes`/`axes.len()` describe a valid slice 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.as_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 @@ -499,6 +560,45 @@ 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 / mean / prod over axis 0. + 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 incompatible_shapes_return_err() { let s = Stream::cpu(); From a666bfe43a70e03a4b9d51b86589823ba4c99c7b Mon Sep 17 00:00:00 2001 From: Kante Yin Date: Tue, 21 Jul 2026 22:44:04 +0100 Subject: [PATCH 3/6] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- crates/mlxcore/src/array.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index 3d03f0a..0094eed 100644 --- a/crates/mlxcore/src/array.rs +++ b/crates/mlxcore/src/array.rs @@ -244,6 +244,9 @@ impl Array { } /// 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) } From cbc137bb9160e81f7609a06df79fca323324993f Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 22:44:19 +0100 Subject: [PATCH 4/6] address comments Signed-off-by: kerthcet --- crates/mlxcore/src/array.rs | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index 0094eed..eb3a469 100644 --- a/crates/mlxcore/src/array.rs +++ b/crates/mlxcore/src/array.rs @@ -329,14 +329,23 @@ impl Array { ) -> 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`/`axes.len()` describe a valid slice for the duration of - // the call; all handles are valid; `op` writes the result into `out`. + // 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.as_ptr(), + axes_ptr, axes.len(), keepdims, stream.as_raw(), From 0bfb86f8fc5c079884f990cbdd48e23bd4019030 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 22:46:23 +0100 Subject: [PATCH 5/6] address comments Signed-off-by: kerthcet --- crates/mlxcore/src/array.rs | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index eb3a469..dba7282 100644 --- a/crates/mlxcore/src/array.rs +++ b/crates/mlxcore/src/array.rs @@ -252,16 +252,25 @@ impl Array { } /// 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) } @@ -522,10 +531,11 @@ mod tests { assert!((logged[0]).abs() < 1e-6); assert!((logged[1] - 1.0).abs() < 1e-6); - // sin(0) == 0, cos(0) == 1. + // 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] From 108e7c1078c77742cfb04a39ff232042a024e7c0 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 21 Jul 2026 22:53:07 +0100 Subject: [PATCH 6/6] address comments Signed-off-by: kerthcet --- crates/mlxcore/src/array.rs | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index dba7282..54f1e1c 100644 --- a/crates/mlxcore/src/array.rs +++ b/crates/mlxcore/src/array.rs @@ -602,7 +602,7 @@ mod tests { let kept = a.sum_axes(&[1], true, &s).unwrap(); assert_eq!(kept.shape(), vec![2, 1]); - // max / min / mean / prod over axis 0. + // 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] @@ -621,6 +621,17 @@ mod tests { ); } + #[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();