feat(api): add fft

Cargo.toml

@@ -21,6 +21,7 @@ num-traits = "0.2.18"
 paste = "1.0.14"
 strum = { version = "0.26", features = ["derive"] }
 thiserror = "1.0.58"
+smallvec = "1"
 
 [dev-dependencies]
 pretty_assertions = "1.4.0"

mlx-macros/Cargo.toml

@@ -17,3 +17,4 @@ proc-macro = true
 [dependencies]
 syn = { version = "2.0.60", features = ["full"] }
 quote = "1.0"
+darling = "0.20"

mlx-macros/src/lib.rs

@@ -1,11 +1,47 @@
 extern crate proc_macro;
+use darling::FromMeta;
 use proc_macro::TokenStream;
 use quote::{format_ident, quote};
 use syn::punctuated::Punctuated;
 use syn::{parse_macro_input, parse_quote, DeriveInput, FnArg, ItemFn, Pat};
 
+#[derive(Debug, FromMeta)]
+enum DeviceType {
+    Cpu,
+    Gpu,
+}
+
+#[derive(Debug)]
+struct DefaultDeviceInput {
+    device: DeviceType,
+}
+
+impl FromMeta for DefaultDeviceInput {
+    fn from_meta(meta: &syn::Meta) -> darling::Result<Self> {
+        let syn::Meta::NameValue(meta_name_value) = meta else {
+            return Err(darling::Error::unsupported_format(
+                "expected a name-value attribute",
+            ));
+        };
+
+        let ident = meta_name_value.path.get_ident().unwrap();
+        assert_eq!(ident, "device", "expected `device`");
+
+        let device = DeviceType::from_expr(&meta_name_value.value)?;
+
+        Ok(DefaultDeviceInput { device })
+    }
+}
+
 #[proc_macro_attribute]
-pub fn default_device(_attr: TokenStream, item: TokenStream) -> TokenStream {
+pub fn default_device(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let input = if !attr.is_empty() {
+        let meta = syn::parse_macro_input!(attr as syn::Meta);
+        Some(DefaultDeviceInput::from_meta(&meta).unwrap())
+    } else {
+        None
+    };
+
     let mut input_fn = parse_macro_input!(item as ItemFn);
     let original_fn = input_fn.clone();
 
@@ -37,8 +73,16 @@ pub fn default_device(_attr: TokenStream, item: TokenStream) -> TokenStream {
     input_fn.sig.inputs = Punctuated::from_iter(filtered_inputs);
 
     // Prepend default stream initialization
-    let default_stream_stmt = parse_quote! {
-        let stream = StreamOrDevice::default();
+    let default_stream_stmt = match input.map(|input| input.device) {
+        Some(DeviceType::Cpu) => parse_quote! {
+            let stream = StreamOrDevice::cpu();
+        },
+        Some(DeviceType::Gpu) => parse_quote! {
+            let stream = StreamOrDevice::gpu();
+        },
+        None => parse_quote! {
+            let stream = StreamOrDevice::default();
+        },
     };
     input_fn.block.stmts.insert(0, default_stream_stmt);
 

src/device.rs

@@ -31,7 +31,8 @@ impl Device {
 
     /// Set the default device.
     ///
-    /// Example:
+    /// # Example:
+    ///
     /// ```rust
     /// use mlx::{Device, DeviceType};
     /// Device::set_default(&Device::new(DeviceType::Cpu, 1));

src/error.rs

@@ -51,3 +51,21 @@ pub enum AsSliceError {
     #[error("dtype mismatch: expected {expecting:?}, found {found:?}")]
     DtypeMismatch { expecting: Dtype, found: Dtype },
 }
+
+#[derive(Error, Debug, PartialEq)]
+pub enum FftError {
+    #[error("fftn requires at least one dimension")]
+    ScalarArray,
+
+    #[error("Invalid axis received for array with {ndim} dimensions")]
+    InvalidAxis { ndim: usize },
+
+    #[error("Shape and axes/axis have different sizes")]
+    IncompatibleShapeAndAxes { shape_size: usize, axes_size: usize },
+
+    #[error("Duplicate axis received: {axis}")]
+    DuplicateAxis { axis: i32 },
+
+    #[error("Invalid output size requested")]
+    InvalidOutputSize,
+}