Struct StateDictMapper 

pub struct StateDictMapper { /* private fields */ }

Implementations

impl StateDictMapper

pub fn new() -> Self

Create a new empty mapper

pub fn rename(self, from: impl Into<String>, to: impl Into<String>) -> Self

Rename a single key

If the key doesn’t exist, this operation is a no-op.

pub fn rename_prefix( self, old_prefix: impl Into<String>, new_prefix: impl Into<String>, ) -> Self

Rename all keys with a given prefix

Example: rename_prefix("old.", "new.") changes “old.weight” to “new.weight”

pub fn strip_prefix(self, prefix: impl Into<String>) -> Self

Strip a prefix from all keys

Example: strip_prefix("model.") changes “model.encoder.weight” to “encoder.weight”

pub fn add_prefix(self, prefix: impl Into<String>) -> Self

Add a prefix to all keys

Example: add_prefix("encoder.") changes “weight” to “encoder.weight”

pub fn transpose(self, key: impl Into<String>) -> Self

Transpose a specific 2D tensor (for Linear weight conversion)

PyTorch Linear layers store weights as [out_features, in_features], while Volta stores them as [in_features, out_features].

This operation only affects 2D tensors. Non-2D tensors are left unchanged.

pub fn transpose_pattern(self, pattern: impl Into<String>) -> Self

Transpose all tensors matching a pattern

Example: transpose_pattern("weight") transposes all keys containing “weight”

pub fn select_keys(self, keys: Vec<String>) -> Self

Select only specific keys (for partial loading)

All other keys are removed from the state dict.

pub fn exclude_keys(self, keys: Vec<String>) -> Self

Exclude specific keys

The specified keys are removed from the state dict.

pub fn transform<F>(self, f: F) -> Self

where F: Fn(&mut StateDict) + 'static,

Apply a custom transformation function

This allows arbitrary transformations beyond the built-in methods.

pub fn apply(&self, state: &mut StateDict)

Apply all transformations to a state dict (in-place)

pub fn map(&self, state: StateDict) -> StateDict

Apply transformations and return new state dict

Trait Implementations

impl Default for StateDictMapper

fn default() -> Self

Returns the “default value” for a type.