ModuleStore

burn_store

Trait ModuleStore 

Source 
pub trait ModuleStore {
    type Error: Debug + Display;

    // Required methods
    fn collect_from<B: Backend, M: ModuleSnapshot<B>>(
        &mut self,
        module: &M,
    ) -> Result<(), Self::Error>;
    fn apply_to<B: Backend, M: ModuleSnapshot<B>>(
        &mut self,
        module: &mut M,
    ) -> Result<ApplyResult, Self::Error>;
    fn get_snapshot(
        &mut self,
        name: &str,
    ) -> Result<Option<&TensorSnapshot>, Self::Error>;
    fn get_all_snapshots(
        &mut self,
    ) -> Result<&BTreeMap<String, TensorSnapshot>, Self::Error>;
    fn keys(&mut self) -> Result<Vec<String>, Self::Error>;
}
Expand description

A trait for handling module storage operations.

ModuleStore provides a unified interface for saving and loading module tensor data with support for various storage formats and advanced features like filtering, remapping, and metadata handling.

Required Associated Types§

Source

type Error: Debug + Display

The error type that can be returned during storage operations.

This should be a format-specific error type that provides detailed information about what went wrong (e.g., I/O errors, format violations, unsupported tensor types).

Required Methods§

Source

fn collect_from<B: Backend, M: ModuleSnapshot<B>>( &mut self, module: &M, ) -> Result<(), Self::Error>

Collect tensor data from a module and store it to storage.

This method traverses the module structure, collects all tensor data according to the store’s configuration (filters, remapping, etc.), and writes it to the underlying storage.

§Arguments
  • module - The module to collect tensor data from. The module must implement ModuleSnapshot to provide tensor access.
§Returns
  • Ok(()) - If all tensors were successfully collected and stored
  • Err(Self::Error) - If an error occurred during collection or writing
Source

fn apply_to<B: Backend, M: ModuleSnapshot<B>>( &mut self, module: &mut M, ) -> Result<ApplyResult, Self::Error>

Load stored tensor data and apply it to a module.

This method reads tensor data from storage and applies it to the provided module. The operation is flexible and can handle partial matches, missing tensors, and extra tensors in the storage.

§Arguments
  • module - The module to apply tensor data to. The module must implement ModuleSnapshot to allow tensor updates.
§Returns
  • Ok(ApplyResult) - Detailed information about the apply operation:
    • applied: List of successfully applied tensor names
    • missing: Tensors expected by the module but not found in storage
    • skipped: Tensors in storage that were not applied (filtered or not needed)
    • errors: Non-critical errors that occurred during apply
  • Err(Self::Error) - If a critical error prevented the apply operation
Source

fn get_snapshot( &mut self, name: &str, ) -> Result<Option<&TensorSnapshot>, Self::Error>

Get a single tensor snapshot by name.

This method provides direct access to individual tensors in storage without requiring a module. The returned TensorSnapshot uses lazy loading - tensor data is only materialized when to_data() is called.

Note: Key remapping is applied, so use the remapped name if configured. Filters are NOT applied - use apply_to() for filtered loading.

Results are cached after the first call for efficient repeated access.

§Arguments
  • name - The tensor name/path (e.g., “encoder.layer1.weight”)
§Returns
  • Ok(Some(&TensorSnapshot)) - Reference to the tensor snapshot if found
  • Ok(None) - If no tensor with that name exists
  • Err(Self::Error) - If an error occurred accessing storage
§Example
ⓘ
let mut store = BurnpackStore::from_file("model.bpk");
if let Some(snapshot) = store.get_snapshot("encoder.weight")? {
    println!("Shape: {:?}", snapshot.shape);
    println!("Dtype: {:?}", snapshot.dtype);
    let data = snapshot.to_data()?;  // Lazy load
}
Source

fn get_all_snapshots( &mut self, ) -> Result<&BTreeMap<String, TensorSnapshot>, Self::Error>

Get all tensor snapshots from storage as an ordered map.

This method returns all tensors in storage as lazy-loading snapshots, organized in a BTreeMap for efficient lookup by name. The map preserves alphabetical ordering of tensor names.

Note: This returns ALL tensors in storage, regardless of any filter settings. Filters are only applied during apply_to(). Key remapping IS applied, so tensor names reflect any configured remapping.

Results are cached after the first call for efficient repeated access.

§Returns
  • Ok(&BTreeMap<String, TensorSnapshot>) - Reference to all tensor snapshots
  • Err(Self::Error) - If an error occurred accessing storage
§Example
ⓘ
let mut store = SafetensorsStore::from_file("model.safetensors");
let snapshots = store.get_all_snapshots()?;
for (name, snapshot) in snapshots {
    println!("{}: {:?}", name, snapshot.shape);
}
Source

fn keys(&mut self) -> Result<Vec<String>, Self::Error>

Get all tensor names/keys in storage.

This method returns the names of all tensors in storage. Useful for inspecting storage contents or checking if specific tensors exist.

Note: Returns ALL tensor names regardless of filter settings. Key remapping IS applied, so names reflect any configured remapping.

§Returns
  • Ok(Vec<String>) - All tensor names in storage
  • Err(Self::Error) - If an error occurred accessing storage
§Example
ⓘ
let mut store = PytorchStore::from_file("model.pth");
let keys = store.keys()?;
println!("Tensors in file: {:?}", keys);

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety", so this trait is not object safe.

Implementors§