pub struct Rope {
pub dims: i32,
pub traditional: bool,
pub base: f32,
pub scale: f32,
/* private fields */
}Expand description
Rotary position embedding of the first dims features (MLX’s nn.RoPE, the fused
fast::rope), optionally with a frequency scaling (Rope::scaled: mlx-lm’s
Llama3RoPE, YarnRoPE, SuScaledRoPE, ProportionalRoPE). No parameters: the
frequencies are computed from the settings, never loaded or saved.
Fields§
§dims: i32Rotated features.
traditional: boolRotate consecutive pairs (true) or the two halves (false, MLX’s default).
base: f32The frequency base (MLX’s default 10000).
scale: f32Scale applied to positions.
Implementations§
Source§impl Rope
impl Rope
Sourcepub fn new(dims: i32) -> Self
pub fn new(dims: i32) -> Self
RoPE over dims features with MLX’s defaults (halves, base 10000, scale 1).
Sourcepub fn scaled(dims: i32, base: f32, scaling: &RopeScaling) -> Self
pub fn scaled(dims: i32, base: f32, scaling: &RopeScaling) -> Self
RoPE over dims features with base and a frequency scaling, as mlx-lm’s scaled
RoPE layers (for RopeScaling::Proportional
dims is the whole head’s).
use nuvec_nn::{Rope, RopeScaling, Yarn};
// gpt-oss: YaRN over 64 features, base 150000, factor 32 over 4096 positions.
let yarn = Yarn { factor: 32.0, original_max_position_embeddings: 4096.0, ..Yarn::default() };
let rope = Rope::scaled(64, 150_000.0, &RopeScaling::Yarn(yarn));
assert!((rope.mscale() - 1.3466).abs() < 1e-4);Sourcepub fn mscale(&self) -> f32
pub fn mscale(&self) -> f32
The input scale the scaling applies to the rotated features (YaRN, LongRoPE; else
1). Callers that rotate queries and keys with Rope::rotation (one fused kernel
each) instead multiply the attention scale by mscale()², which is the same up to
rounding when the rotated features are the whole head. In half precision mlx-lm
multiplies q and k by the mscale rounded to their dtype, which the fold does not
reproduce: nuvec-lm multiplies explicitly there and folds only in float32.
Sourcepub fn rotation(&self) -> Rotation<'_>
pub fn rotation(&self) -> Rotation<'_>
The rotation for the fused ops (fast::rope, fast::norm_rope), without the input
scale (Rope::mscale).
Sourcepub fn traditional(self) -> Self
pub fn traditional(self) -> Self
Rotates consecutive pairs of features (MLX’s traditional=True).
Trait Implementations§
Source§impl Tree for Ropewhere
Self: Clone,
impl Tree for Ropewhere
Self: Clone,
Source§fn visit_mut(&mut self, __v: &mut VisitMut<'_>)
fn visit_mut(&mut self, __v: &mut VisitMut<'_>)
v mutably, in the order of
Leaves::visit.Source§fn map(&self, f: impl FnMut(&str, &Array) -> Array) -> Self
fn map(&self, f: impl FnMut(&str, &Array) -> Array) -> Self
f(path, leaf) (same structure; parameters keep
their frozen flag).Source§fn flatten(&self) -> Vec<(String, Array)>
fn flatten(&self) -> Vec<(String, Array)>
tree_flatten).Source§fn unflatten(
&self,
leaves: impl IntoIterator<Item = Array>,
) -> Result<Self, Error>
fn unflatten( &self, leaves: impl IntoIterator<Item = Array>, ) -> Result<Self, Error>
self is the template) with leaves in order: exactly
Tree::len of them, else an error naming both counts.Source§fn load_flat(
&mut self,
named: impl IntoIterator<Item = (String, Array)>,
mode: LoadMode,
) -> Result<(), Error>
fn load_flat( &mut self, named: impl IntoIterator<Item = (String, Array)>, mode: LoadMode, ) -> Result<(), Error>
LoadMode::Strict requires exactly the tree’s paths with
matching shapes; LoadMode::Partial assigns the matching names and ignores
unknown ones (shapes must still match). Loaded dtypes are kept. On any problem the
error lists every one (missing, unknown and mismatched names, sorted) and the tree is
unchanged. Duplicate paths (in the tree or in named) are an error.Auto Trait Implementations§
impl !RefUnwindSafe for Rope
impl !UnwindSafe for Rope
impl Freeze for Rope
impl Send for Rope
impl Sync for Rope
impl Unpin for Rope
impl UnsafeUnpin for Rope
Blanket Implementations§
impl<T> AutoreleaseSafe for Twhere
T: ?Sized,
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
impl<ST, DT> CastableFrom<ST, Initialized, Initialized> for DT
impl<ST, DT> CastableFrom<ST, Uninit, Uninit> for DT
Source§impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> CloneToUninit for Twhere
T: Clone,
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§impl<T> IntoResult for Twhere
T: Tree,
impl<T> IntoResult for Twhere
T: Tree,
Source§impl<T> Module for Twhere
T: Tree,
impl<T> Module for Twhere
T: Tree,
Source§fn parameters(&self) -> Vec<(String, Array)>
fn parameters(&self) -> Vec<(String, Array)>
Source§fn trainable_parameters(&self) -> Vec<(String, Array)>
fn trainable_parameters(&self) -> Vec<(String, Array)>
Source§fn num_params(&self) -> usize
fn num_params(&self) -> usize
tree_flatten total).Source§fn num_trainable(&self) -> usize
fn num_trainable(&self) -> usize
Source§fn summary(&self) -> Summary
fn summary(&self) -> Summary
Source§fn display(&self) -> ModuleDisplay<'_, Self>
fn display(&self) -> ModuleDisplay<'_, Self>
println!("{}", model.display()).Source§fn freeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize>
fn freeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize>
filter, returning how many changed.
A matching plain Array leaf cannot be frozen (only Params can): that is an
error naming the paths, and nothing changes. Buffers are never trainable anyway.Source§fn unfreeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize>
fn unfreeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize>
filter, returning how many changed.
Buffers (quantized weights) and integer parameters stay non-trainable.Source§fn train(&mut self, on: bool)
fn train(&mut self, on: bool)
true, the default) or evaluation
(false, where dropout is the identity), as MLX’s Module.train(mode). (There is
no eval(): it would read like x.eval(), which evaluates.)Source§fn training(&self) -> bool
fn training(&self) -> bool
true without any).Source§fn load_weights(&mut self, src: impl WeightSource) -> Result<()>
fn load_weights(&mut self, src: impl WeightSource) -> Result<()>
load_weights(strict=True)): every leaf
must be given with its shape and nothing else (loaded dtypes are kept). On any
problem the error lists every missing, unknown and mismatched name, and the model
is unchanged. src is a safetensors path or named arrays.Source§fn load_weights_partial(&mut self, src: impl WeightSource) -> Result<()>
fn load_weights_partial(&mut self, src: impl WeightSource) -> Result<()>
load_weights(strict=False)): unknown
names are ignored, missing leaves keep their values, shapes must match.Source§fn save_weights(&self, path: impl AsRef<Path>) -> Result<()>
fn save_weights(&self, path: impl AsRef<Path>) -> Result<()>
save_weights), evaluating them. Unloaded placeholders are an error naming them.Source§fn save_trainable(&self, path: impl AsRef<Path>) -> Result<()>
fn save_trainable(&self, path: impl AsRef<Path>) -> Result<()>
Source§fn map_modules<T: Tree>(
&mut self,
filter: impl Fn(&str) -> bool,
f: impl FnMut(&str, &mut T) -> Result<()>,
) -> Result<usize>
fn map_modules<T: Tree>( &mut self, filter: impl Fn(&str) -> bool, f: impl FnMut(&str, &mut T) -> Result<()>, ) -> Result<usize>
f(path, submodule) on every submodule of type T whose path matches
filter (f may replace it, as when applying LoRA), returning how many. It stops at
the first error from f; submodules changed before it keep their change. Matching
submodules are not searched for further matches.