address comments

Author: LYZJU2019

mlx-lm/Cargo.toml

@@ -13,6 +13,7 @@ description = "Rust implementation of mlx-lm"
 [dependencies]
 # Local dependencies
 mlx-rs.workspace = true
+mlx-macros.workspace = true
 mlx-lm-utils.workspace = true
 
 # External dependencies

mlx-lm/src/utils/rope.rs

@@ -1,6 +1,14 @@
 use std::collections::HashMap;
 
-use mlx_rs::{builder::Builder, error::Exception, module::Module, nn, Array};
+use mlx_macros::ModuleParameters;
+use mlx_rs::{
+    builder::Builder,
+    error::Exception,
+    module::Module,
+    nn,
+    ops::{arange, which},
+    Array,
+};
 use serde::Deserialize;
 
 #[derive(Debug, Clone, PartialEq)]
@@ -9,7 +17,7 @@ pub enum FloatOrStr<'a> {
     Str(&'a str),
 }
 
-// TODO: check if additionl serde attributes are needed
+// TODO: check if additional serde attributes are needed
 #[derive(Debug, Clone, Deserialize)]
 #[serde(untagged)]
 pub enum FloatOrString {
@@ -26,7 +34,12 @@ impl FloatOrString {
     }
 }
 
-fn get_float_from_config(
+/// Get a numeric float value from a scaling config by key.
+///
+/// Note: str variants in the config are not always floats — values like "default" or "linear"
+/// are also valid for non-numeric fields. This function should only be called for keys that
+/// are expected to hold numeric values.
+fn get_numeric_from_config(
     config: &HashMap<String, FloatOrString>,
     key: &str,
 ) -> Result<f32, Exception> {
@@ -39,19 +52,21 @@ fn get_float_from_config(
         FloatOrStr::Float(f) => Ok(f),
         FloatOrStr::Str(s) => s
             .parse::<f32>()
-            .map_err(|_| Exception::custom(format!(r#"key "{key}" is not a valid float"#))),
+            .map_err(|_| Exception::custom(format!(r#"key "{key}" is not a valid number"#))),
     }
 }
 
 /// Llama3-style RoPE with frequency scaling.
 ///
 /// Applies piecewise frequency scaling based on wavelength cutoffs derived from
 /// `low_freq_factor`, `high_freq_factor`, `factor`, and `original_max_position_embeddings`.
-#[derive(Debug, Clone)]
+// TODO: support derive ModuleParameters for structs with non-param Array fields
+#[derive(Debug, Clone, ModuleParameters)]
 pub struct Llama3Rope {
     pub dimensions: i32,
     pub traditional: bool,
     pub scale: f32,
+    /// Pre-computed scaled frequencies. Not a module parameter.
     pub freqs: Array,
 }
 
@@ -67,12 +82,12 @@ impl Llama3Rope {
     ) -> Result<Self, Exception> {
         let half_dims = dims / 2;
 
-        // Compute freqs as periods: base^(2i/dims), matching Python:
+        // Compute freqs using MLX ops, matching Python:
         //   freqs = base ** (mx.arange(0, dims, 2) / dims)
-        let mut freqs = Vec::with_capacity(half_dims as usize);
-        for i in 0..half_dims {
-            freqs.push(base.powf(2.0 * i as f32 / dims as f32));
-        }
+        // which equals base^(2i/dims) for i in 0..half_dims
+        let indices = arange::<_, f32>(None, half_dims, None)?;
+        let exponents = indices.multiply(Array::from_f32(2.0 / dims as f32))?;
+        let freqs = Array::from_f32(base).power(&exponents)?;
 
         let old_context_len = original_max_position_embeddings as f32;
         let low_freq_wavelen = old_context_len / low_freq_factor;
@@ -85,68 +100,42 @@ impl Llama3Rope {
         //   smooth_factors = (old_context_len / wavelens - low_freq_factor) / (high - low)
         //   smooth_freqs = freqs / ((1 - smooth_factors) / factor + smooth_factors)
         //   freqs = where(is_medium, smooth_freqs, freqs)
-        let mut scaled_freqs = Vec::with_capacity(half_dims as usize);
-        for &freq in &freqs {
-            let wavelen = 2.0 * std::f32::consts::PI * freq;
-            // First pass: scale low frequencies (long wavelengths) by factor
-            let freq = if wavelen > low_freq_wavelen {
-                freq * factor
-            } else {
-                freq
-            };
-            // Second pass: apply smooth interpolation for medium frequencies
-            let is_medium = wavelen > high_freq_wavelen && wavelen < low_freq_wavelen;
-            if is_medium {
-                let smooth_factor = (old_context_len / wavelen - low_freq_factor)
-                    / (high_freq_factor - low_freq_factor);
-                let smooth_freq = freq / ((1.0 - smooth_factor) / factor + smooth_factor);
-                scaled_freqs.push(smooth_freq);
-            } else {
-                scaled_freqs.push(freq);
-            }
-        }
+        let two_pi = Array::from_f32(2.0 * std::f32::consts::PI);
+        let wavelens = freqs.multiply(&two_pi)?;
 
-        let freqs_array = Array::from_slice(&scaled_freqs, &[half_dims]);
+        // First pass: scale low frequencies (long wavelengths) by factor
+        let is_low = wavelens.gt(Array::from_f32(low_freq_wavelen))?;
+        let freqs = which(&is_low, &freqs.multiply(Array::from_f32(factor))?, &freqs)?;
+
+        // Second pass: smooth interpolation for medium frequencies
+        let is_medium = wavelens
+            .gt(Array::from_f32(high_freq_wavelen))?
+            .logical_and(&wavelens.lt(Array::from_f32(low_freq_wavelen))?)?;
+
+        let smooth_factors = wavelens
+            .reciprocal()?
+            .multiply(Array::from_f32(old_context_len))?
+            .subtract(Array::from_f32(low_freq_factor))?
+            .divide(Array::from_f32(high_freq_factor - low_freq_factor))?;
+
+        // smooth_freqs = freqs / ((1 - smooth_factors) / factor + smooth_factors)
+        let one_minus_smooth = Array::from_f32(1.0).subtract(&smooth_factors)?;
+        let denom = one_minus_smooth
+            .divide(Array::from_f32(factor))?
+            .add(&smooth_factors)?;
+        let smooth_freqs = freqs.divide(&denom)?;
+
+        let freqs = which(&is_medium, &smooth_freqs, &freqs)?;
 
         Ok(Self {
             dimensions: dims,
             traditional,
             scale: 1.0,
-            freqs: freqs_array,
+            freqs,
         })
     }
 }
 
-impl mlx_rs::module::ModuleParameters for Llama3Rope {
-    fn num_parameters(&self) -> usize {
-        0
-    }
-
-    fn freeze_parameters(&mut self, _recursive: bool) {}
-
-    fn unfreeze_parameters(&mut self, _recursive: bool) {}
-
-    fn parameters(&self) -> mlx_rs::module::ModuleParamRef<'_> {
-        mlx_rs::nested::NestedHashMap::new()
-    }
-
-    fn parameters_mut(&mut self) -> mlx_rs::module::ModuleParamMut<'_> {
-        mlx_rs::nested::NestedHashMap::new()
-    }
-
-    fn trainable_parameters(&self) -> mlx_rs::module::ModuleParamRef<'_> {
-        mlx_rs::nested::NestedHashMap::new()
-    }
-
-    fn all_frozen(&self) -> Option<bool> {
-        None
-    }
-
-    fn any_frozen(&self) -> Option<bool> {
-        None
-    }
-}
-
 impl<'a, Input> Module<Input> for Llama3Rope
 where
     Input: Into<nn::RopeInput<'a>>,
@@ -181,6 +170,7 @@ pub enum RopeVariant {
     Llama3(Llama3Rope),
 }
 
+// TODO: support derive ModuleParameters for enum
 impl mlx_rs::module::ModuleParameters for RopeVariant {
     fn num_parameters(&self) -> usize {
         0
@@ -256,8 +246,7 @@ pub fn initialize_rope(
 
     if rope_type == FloatOrStr::Str("default") || rope_type == FloatOrStr::Str("linear") {
         let scale = if rope_type == FloatOrStr::Str("linear") {
-            let den = get_float_from_config(scaling_config.as_ref().unwrap(), "factor")?;
-
+            let den = get_numeric_from_config(scaling_config.as_ref().unwrap(), "factor")?;
             1.0 / den
         } else {
             1.0
@@ -275,11 +264,11 @@ pub fn initialize_rope(
             .as_ref()
             .ok_or_else(|| Exception::custom("scaling_config is required for llama3 RoPE"))?;
 
-        let factor = get_float_from_config(config, "factor")?;
-        let low_freq_factor = get_float_from_config(config, "low_freq_factor")?;
-        let high_freq_factor = get_float_from_config(config, "high_freq_factor")?;
+        let factor = get_numeric_from_config(config, "factor")?;
+        let low_freq_factor = get_numeric_from_config(config, "low_freq_factor")?;
+        let high_freq_factor = get_numeric_from_config(config, "high_freq_factor")?;
         let original_max_position_embeddings =
-            get_float_from_config(config, "original_max_position_embeddings")? as i32;
+            get_numeric_from_config(config, "original_max_position_embeddings")? as i32;
 
         let rope = Llama3Rope::new(
             dims,