Use custom c shim for fallible closure

mlx-rs/examples/tutorial.rs

@@ -83,10 +83,10 @@ fn automatic_differentiation() {
 
     let x = Array::from(1.5);
 
-    let mut dfdx = calculate_grad(f, &x);
+    let dfdx = calculate_grad(f, &x);
     assert_eq!(dfdx.item::<f32>(), 2.0 * 1.5);
 
-    let mut dfdx2 = calculate_grad(|args| calculate_grad(f, args), &x);
+    let dfdx2 = calculate_grad(|args| calculate_grad(f, args), &x);
     assert_eq!(dfdx2.item::<f32>(), 2.0);
 }
 

mlx-rs/src/array/operators.rs

@@ -1,4 +1,4 @@
-use crate::{prelude::ScalarOrArray, Array, StreamOrDevice};
+use crate::{utils::ScalarOrArray, Array, StreamOrDevice};
 use num_traits::Pow;
 use std::ops::{
     Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Not, Rem, RemAssign, Sub, SubAssign,
@@ -14,7 +14,7 @@ macro_rules! impl_binary_op {
 
             fn $method(self, rhs: T) -> Self::Output {
                 paste::paste! {
-                    self.[<$c_method _device>](rhs, StreamOrDevice::default()).unwrap()
+                    self.[<$c_method _device>](rhs.into_owned_or_ref_array(), StreamOrDevice::default()).unwrap()
                 }
             }
         }
@@ -27,7 +27,7 @@ macro_rules! impl_binary_op {
 
             fn $method(self, rhs: T) -> Self::Output {
                 paste::paste! {
-                    self.[<$c_method _device>](rhs, StreamOrDevice::default()).unwrap()
+                    self.[<$c_method _device>](rhs.into_owned_or_ref_array(), StreamOrDevice::default()).unwrap()
                 }
             }
         }

mlx-rs/src/error.rs

@@ -43,8 +43,8 @@ impl Exception {
 }
 
 thread_local! {
-    pub static LAST_MLX_ERROR: Cell<*const c_char> = const { Cell::new(std::ptr::null()) };
-    pub static IS_HANDLER_SET: Cell<bool> = const { Cell::new(false) };
+    static LAST_MLX_ERROR: Cell<*const c_char> = const { Cell::new(std::ptr::null()) };
+    static IS_HANDLER_SET: Cell<bool> = const { Cell::new(false) };
 }
 
 #[no_mangle]

mlx-rs/src/lib.rs

@@ -28,7 +28,6 @@ pub mod prelude {
         dtype::Dtype,
         ops::indexing::{Ellipsis, IndexMutOp, IndexOp, IntoStrideBy, NewAxis},
         stream::StreamOrDevice,
-        utils::ScalarOrArray,
     };
 
     pub use num_traits::Pow;

mlx-rs/src/macros/array.rs

@@ -41,9 +41,19 @@
 ///         [10, 11, 12]
 ///     ]
 /// ]);
+///
+/// // Create a 2x2 array by specifying the shape
+/// let a = array!([1, 2, 3, 4], shape=[2, 2]);
 /// ```
 #[macro_export]
 macro_rules! array {
+    ([$($x:expr),*], shape=[$($s:expr),*]) => {
+        {
+            let data = [$($x,)*];
+            let shape = [$($s,)*];
+            $crate::Array::from_slice(&data, &shape)
+        }
+    };
     ([$([$([$($x:expr),*]),*]),*]) => {
         {
             let arr = [$([$([$($x,)*],)*],)*];
@@ -133,4 +143,16 @@ mod tests {
         assert_eq!(a.index((1, 1, 1)).item::<i32>(), 11);
         assert_eq!(a.index((1, 1, 2)).item::<i32>(), 12);
     }
+
+    #[test]
+    fn test_array_with_shape() {
+        let a = array!([1, 2, 3, 4], shape = [2, 2]);
+
+        assert!(a.ndim() == 2);
+        assert_eq!(a.shape(), &[2, 2]);
+        assert_eq!(a.index((0, 0)).item::<i32>(), 1);
+        assert_eq!(a.index((0, 1)).item::<i32>(), 2);
+        assert_eq!(a.index((1, 0)).item::<i32>(), 3);
+        assert_eq!(a.index((1, 1)).item::<i32>(), 4);
+    }
 }

mlx-rs/src/macros/internal.rs

@@ -8,12 +8,33 @@ macro_rules! try_catch_c_ptr_expr {
         if c_ptr.is_null() {
             // SAFETY: there must be an error if the pointer is null
             return Err($crate::error::get_and_clear_last_mlx_error()
+                // .or($crate::error::take_last_mlx_closure_error())
                 .expect("A null pointer was returned, but no error was set."));
         }
         c_ptr
     }};
 }
 
+macro_rules! try_catch_mlx_closure_error {
+    ($expr:expr) => {{
+        if !$crate::error::is_mlx_error_handler_set() {
+            $crate::error::setup_mlx_error_handler();
+        }
+
+        let c_ptr = $expr;
+        // Always check for closure errors
+        if let Some(error) = $crate::error::get_and_clear_last_mlx_error()
+        // .or($crate::error::take_last_mlx_closure_error())
+        {
+            return Err(error);
+        }
+        if c_ptr.is_null() {
+            panic!("A null pointer was returned, but no error was set.");
+        }
+        c_ptr
+    }};
+}
+
 /// See `assertEqual` in the swift binding tests
 #[allow(unused_macros)]
 macro_rules! assert_array_all_close {