Merge pull request #272 from oxideai/release/0.25.1

release: Version 0.25.1

Author: minghuaw

Cargo.toml

@@ -1,7 +1,7 @@
 [workspace.package]
 # All but mlx-sys should follow the same version. mlx-sys should follow 
 # the version of mlx-c.
-version = "0.25.0"
+version = "0.25.1"
 edition = "2021"
 authors = [
     "Minghua Wu <michael.wu1107@gmail.com>",
@@ -29,9 +29,9 @@ resolver = "2"
 [workspace.dependencies]
 # workspace local dependencies
 mlx-sys = { version = "=0.2.0", path = "mlx-sys" }
-mlx-macros = { version = "0.25.0", path = "mlx-macros" }
-mlx-internal-macros = { version = "0.25.0", path = "mlx-internal-macros" }
-mlx-rs = { version = "0.25.0", path = "mlx-rs" }
+mlx-macros = { version = "0.25", path = "mlx-macros" }
+mlx-internal-macros = { version = "0.25", path = "mlx-internal-macros" }
+mlx-rs = { version = "0.25.1", path = "mlx-rs" }
 
 # external dependencies
 thiserror = "2"

mlx-rs/CHANGELOG.md

@@ -1,8 +1,12 @@
 # Changelog
 
-## 0.25.0-alpha.1
+## 0.25.1
 
-- Update `mlx-c` to version "0.2.0-alpha" and changes function signatures to
+- Fix bug with `index_mut`
+
+## 0.25.0
+
+- Update `mlx-c` to version "0.2.0" and changes function signatures to
   match the new API
 - Update `thiserror` to version "2"
 - Fix wrong states number in `compile_with_state`

mlx-rs/src/ops/arithmetic.rs

@@ -652,6 +652,22 @@ impl Array {
             mlx_sys::mlx_square(res, self.as_ptr(), stream.as_ref().as_ptr())
         })
     }
+
+    /// Element-wise real part from a complex array.
+    #[default_device]
+    pub fn real_device(&self, stream: impl AsRef<Stream>) -> Result<Array> {
+        Array::try_from_op(|res| unsafe {
+            mlx_sys::mlx_real(res, self.as_ptr(), stream.as_ref().as_ptr())
+        })
+    }
+
+    /// Element-wise imag part from a complex array.
+    #[default_device]
+    pub fn imag_device(&self, stream: impl AsRef<Stream>) -> Result<Array> {
+        Array::try_from_op(|res| unsafe {
+            mlx_sys::mlx_imag(res, self.as_ptr(), stream.as_ref().as_ptr())
+        })
+    }
 }
 
 /// Element-wise absolute value.
@@ -1306,6 +1322,24 @@ pub fn tanh_device(a: impl AsRef<Array>, #[optional] stream: impl AsRef<Stream>)
     })
 }
 
+/// Element-wise real part from a complex array.
+#[generate_macro]
+#[default_device]
+pub fn real_device(a: impl AsRef<Array>, #[optional] stream: impl AsRef<Stream>) -> Result<Array> {
+    Array::try_from_op(|res| unsafe {
+        mlx_sys::mlx_real(res, a.as_ref().as_ptr(), stream.as_ref().as_ptr())
+    })
+}
+
+/// Element-wise imaginary part from a complex array.
+#[generate_macro]
+#[default_device]
+pub fn imag_device(a: impl AsRef<Array>, #[optional] stream: impl AsRef<Stream>) -> Result<Array> {
+    Array::try_from_op(|res| unsafe {
+        mlx_sys::mlx_imag(res, a.as_ref().as_ptr(), stream.as_ref().as_ptr())
+    })
+}
+
 /// Matrix multiplication with block masking.
 ///
 /// See the [python API docs](
@@ -2500,6 +2534,13 @@ mod tests {
         .item::<bool>());
     }
 
+    #[test]
+    fn test_unary_real_imag() {
+        let x = Array::from_complex(complex64::new(0.0, 1.0));
+        assert_eq!(real(&x).unwrap(), Array::from_f32(0.0));
+        assert_eq!(imag(&x).unwrap(), Array::from_f32(1.0));
+    }
+
     #[test]
     fn test_binary_add() {
         let x = array![1.0];

mlx-rs/src/ops/indexing/indexmut_impl.rs

@@ -8,6 +8,7 @@ use crate::{
     ops::{
         broadcast_arrays_device, broadcast_to_device,
         indexing::{count_non_new_axis_operations, expand_ellipsis_operations},
+        reshape_device,
     },
     utils::{resolve_index_signed_unchecked, VectorArray},
     Array, Stream,
@@ -84,40 +85,59 @@ fn update_slice(
     let operations = expand_ellipsis_operations(ndim, operations);
 
     // If no non-None indices return the broadcasted update
-    if count_non_new_axis_operations(&operations) == 0 {
+    let non_new_axis_operation_count = count_non_new_axis_operations(&operations);
+    if non_new_axis_operation_count == 0 {
         return Ok(Some(broadcast_to_device(&update, src.shape(), &stream)?));
     }
 
     // Process entries
-    let mut update_expand_dims: SmallVec<[i32; DEFAULT_STACK_VEC_LEN]> = SmallVec::new();
-    let mut axis = 0i32;
-    for item in operations.iter() {
+    // let mut update_expand_dims: SmallVec<[i32; DEFAULT_STACK_VEC_LEN]> = SmallVec::new();
+    let mut update_reshape: SmallVec<[i32; DEFAULT_STACK_VEC_LEN]> = smallvec![0; ndim];
+    let mut axis = src.ndim() - 1;
+    let mut update_axis = update.ndim() as i32 - 1;
+
+    while axis >= non_new_axis_operation_count {
+        if update_axis >= 0 {
+            update_reshape[axis] = update.dim(update_axis);
+            update_axis -= 1;
+        } else {
+            update_reshape[axis] = 1;
+        }
+        axis -= 1;
+    }
+
+    for item in operations.iter().rev() {
         use ArrayIndexOp::*;
 
         match item {
             TakeIndex { index } => {
-                let size = src.dim(axis);
+                let size = src.dim(axis as i32);
                 let index = if index.is_negative() {
                     size + index
                 } else {
                     *index
                 };
                 // SAFETY: axis is always non-negative
-                starts[axis as usize] = index;
-                ends[axis as usize] = index + 1;
-                if ndim - (axis as usize) < update.ndim() {
-                    update_expand_dims.push(axis.saturating_sub_unsigned(ndim as u32));
-                }
+                starts[axis] = index;
+                ends[axis] = index.saturating_add(1);
 
-                axis = axis.saturating_add(1);
+                update_reshape[axis] = 1;
+                axis = axis.saturating_sub(1);
             }
-            Slice(range_index) => {
-                let size = src.dim(axis);
+            Slice(slice) => {
+                let size = src.dim(axis as i32);
                 // SAFETY: axis is always non-negative
-                starts[axis as usize] = range_index.start(size);
-                ends[axis as usize] = range_index.end(size);
-                strides[axis as usize] = range_index.stride();
-                axis = axis.saturating_add(1);
+                starts[axis] = slice.start(size);
+                ends[axis] = slice.end(size);
+                strides[axis] = slice.stride();
+
+                if update_axis >= 0 {
+                    update_reshape[axis] = update.dim(update_axis);
+                    update_axis = update_axis.saturating_sub(1);
+                } else {
+                    update_reshape[axis] = 1;
+                }
+                axis = axis.saturating_sub(1);
             }
             ExpandDims => {}
             Ellipsis | TakeArray { indices: _ } | TakeArrayRef { indices: _ } => {
@@ -126,8 +146,8 @@ fn update_slice(
         }
     }
 
-    if !update_expand_dims.is_empty() {
-        update = Cow::Owned(update.expand_dims_axes_device(&update_expand_dims, &stream)?);
+    if update.shape() != &update_reshape[..] {
+        update = Cow::Owned(reshape_device(update, &update_reshape, &stream)?);
     }
 
     Ok(Some(src.slice_update_device(
@@ -1277,7 +1297,10 @@ where
 /// The unit tests below are adapted from the Swift binding tests
 #[cfg(test)]
 mod tests {
-    use crate::{ops::indexing::*, Array};
+    use crate::{
+        ops::{indexing::*, ones, zeros},
+        Array,
+    };
 
     #[test]
     fn test_array_mutate_single_index() {
@@ -1473,4 +1496,16 @@ mod tests {
             128142
         );
     }
+
+    #[test]
+    fn test_slice_update_with_broadcast() {
+        let mut xs = zeros::<f32>(&[4, 3, 2]).unwrap();
+        let x = ones::<f32>(&[4, 2]).unwrap();
+
+        let result = xs.try_index_mut((.., 0, ..), x);
+        assert!(
+            result.is_ok(),
+            "Failed to update slice with broadcast: {result:?}"
+        );
+    }
 }