Skip to main content

OptimizerState

mlx_rs::optimizers

Trait OptimizerState 

Source 
pub trait OptimizerState: Sized {
    type UnflattenError: Error + Into<IoError>;

    // Required methods
    fn state_projection(
        &mut self,
    ) -> Result<StateProjection<'_>, StateProjectionError>;
    fn unflatten<I, K>(input: I) -> Result<Self, Self::UnflattenError>
       where I: IntoIterator<Item = (K, Array)>,
             K: Ord + AsRef<str> + Into<Rc<str>>;

    // Provided methods
    fn flatten(
        &mut self,
    ) -> Result<Vec<(Rc<str>, &Array)>, StateProjectionError> { ... }
    fn flatten_mut(
        &mut self,
    ) -> Result<Vec<(Rc<str>, &mut Array)>, StateProjectionError> { ... }
    fn save_safetensors(
        &mut self,
        path: impl AsRef<Path>,
    ) -> Result<(), IoError> { ... }
    fn load_safetensors(
        &mut self,
        path: impl AsRef<Path>,
    ) -> Result<(), IoError> { ... }
}
Expand description

Trait for optimizer states.

Required Associated Types§

Source

type UnflattenError: Error + Into<IoError>

Error type for unflatten.

Required Methods§

Source

fn state_projection( &mut self, ) -> Result<StateProjection<'_>, StateProjectionError>

Declare all required and optional optimizer-state slots.

Source

fn unflatten<I, K>(input: I) -> Result<Self, Self::UnflattenError>
where I: IntoIterator<Item = (K, Array)>, K: Ord + AsRef<str> + Into<Rc<str>>,

Unflatten an iterator of key-value pairs into the optimizer state.

Provided Methods§

Source

fn flatten(&mut self) -> Result<Vec<(Rc<str>, &Array)>, StateProjectionError>

Flatten present state entries in stable key order.

Source

fn flatten_mut( &mut self, ) -> Result<Vec<(Rc<str>, &mut Array)>, StateProjectionError>

Flatten mutable present state entries in stable key order.

Source

fn save_safetensors(&mut self, path: impl AsRef<Path>) -> Result<(), IoError>

Save the optimizer state to a safetensors file.

Source

fn load_safetensors(&mut self, path: impl AsRef<Path>) -> Result<(), IoError>

Load the optimizer state from a safetensors file.

Dyn Compatibility§

This trait is not dyn compatible.

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

Implementors§