Trait Module 

pub trait Module: Tree {
    // Provided methods
    fn parameters(&self) -> Vec<(String, Array)> { ... }
    fn trainable_parameters(&self) -> Vec<(String, Array)> { ... }
    fn num_params(&self) -> usize { ... }
    fn num_trainable(&self) -> usize { ... }
    fn summary(&self) -> Summary { ... }
    fn display(&self) -> ModuleDisplay<'_, Self> { ... }
    fn freeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize> { ... }
    fn unfreeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize> { ... }
    fn train(&mut self, on: bool) { ... }
    fn training(&self) -> bool { ... }
    fn load_weights(&mut self, src: impl WeightSource) -> Result<()> { ... }
    fn load_weights_partial(&mut self, src: impl WeightSource) -> Result<()> { ... }
    fn save_weights(&self, path: impl AsRef<Path>) -> Result<()> { ... }
    fn save_trainable(&self, path: impl AsRef<Path>) -> Result<()> { ... }
    fn map_modules<T: Tree>(
        &mut self,
        filter: impl Fn(&str) -> bool,
        f: impl FnMut(&str, &mut T) -> Result<()>,
    ) -> Result<usize> { ... }
}

Model operations, as MLX’s nn.Module, for every Tree (a blanket impl): a model is a struct of layers with #[derive(Module)], and its leaves’ paths are MLX’s names.

Leaves are parameters (Param, trainable unless frozen), buffers (never trained, such as quantized weights) or plain arrays.

Provided Methods

fn parameters(&self) -> Vec<(String, Array)>

Every leaf (parameters, buffers and plain arrays) with its path, in order.

fn trainable_parameters(&self) -> Vec<(String, Array)>

The trainable leaves with their paths, in order.

fn num_params(&self) -> usize

The number of elements over every leaf (MLX’s tree_flatten total).

fn num_trainable(&self) -> usize

The number of elements over the trainable leaves.

fn summary(&self) -> Summary

A table of the leaves (path, shape, dtype, trainable, bytes) with totals.

fn display(&self) -> ModuleDisplay<'_, Self>

The nested modules with their parameter counts, for printing: println!("{}", model.display()).

fn freeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize>

Freezes the parameters whose path matches 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.

fn unfreeze(&mut self, filter: impl Fn(&str) -> bool) -> Result<usize>

Unfreezes the parameters whose path matches filter, returning how many changed. Buffers (quantized weights) and integer parameters stay non-trainable.

fn train(&mut self, on: bool)

Sets every mode flag of the tree: training (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.)

fn training(&self) -> bool

Whether the tree is in training mode (its first mode flag; true without any).

fn load_weights(&mut self, src: impl WeightSource) -> Result<()>

Replaces leaves by name, strictly (MLX’s 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.

fn load_weights_partial(&mut self, src: impl WeightSource) -> Result<()>

Replaces the leaves given by name (MLX’s load_weights(strict=False)): unknown names are ignored, missing leaves keep their values, shapes must match.

fn save_weights(&self, path: impl AsRef<Path>) -> Result<()>

Writes every leaf as safetensors with MLX’s names and dtypes (MLX’s save_weights), evaluating them. Unloaded placeholders are an error naming them.

fn save_trainable(&self, path: impl AsRef<Path>) -> Result<()>

Writes only the trainable leaves (an adapter’s weights, say).

fn map_modules<T: Tree>( &mut self, filter: impl Fn(&str) -> bool, f: impl FnMut(&str, &mut T) -> Result<()>, ) -> Result<usize>

Calls 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.

Dyn Compatibility

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors

impl<T: Tree> Module for T