Trait WeightLoader

Source

pub trait WeightLoader: Send {
    // Required methods
    fn len(&self) -> usize;
    fn take(&mut self, key: &str) -> Result<(Vec<f32>, Vec<usize>), Error>;
    fn take_transposed(
        &mut self,
        key: &str,
    ) -> Result<(Vec<f32>, Vec<usize>), Error>;
    fn remaining_keys(&self) -> Vec<String>;

    // Provided methods
    fn format_id(&self) -> &'static str { ... }
    fn is_empty(&self) -> bool { ... }
    fn take_packed(
        &mut self,
        key: &str,
    ) -> Result<Option<(Vec<u8>, QuantScheme, Vec<usize>)>, Error> { ... }
    fn tensor_bytes_borrowed(&self, key: &str) -> Option<&[u8]> { ... }
    fn arch_hint(&self) -> Option<&str> { ... }
}

Expand description

Common interface every weight format must satisfy. Mirrors the existing WeightMap API so the safetensors impl is a one-line adapter.

Register additional on-disk formats with crate::weight_registry::register_weight_format.

Required Methods§

Source

fn len(&self) -> usize

Number of distinct weights in the file.

Source

fn take(&mut self, key: &str) -> Result<(Vec<f32>, Vec<usize>), Error>

Take the named tensor as (f32_data, shape). Removes from the loader so callers can detect “weights I never used.”

Source

fn take_transposed( &mut self, key: &str, ) -> Result<(Vec<f32>, Vec<usize>), Error>

Same as take but transposed (last two dims swapped). Most safetensors weights are stored row-major-of-PyTorch convention, which RLX’s IR consumes column-major; this helper is the convention-bridge.

Source

fn remaining_keys(&self) -> Vec<String>

Names that haven’t been taken yet — useful for “did the model use every weight?” hygiene checks.

Provided Methods§

Source

fn format_id(&self) -> &'static str

Format id (safetensors, gguf, or a custom registration).

Source

fn is_empty(&self) -> bool

Source

fn take_packed( &mut self, key: &str, ) -> Result<Option<(Vec<u8>, QuantScheme, Vec<usize>)>, Error>

Take packed K-quant bytes when supported; default returns None.

Source

fn tensor_bytes_borrowed(&self, key: &str) -> Option<&[u8]>

Borrow packed bytes without marking taken (GGUF mmap path).

Source

fn arch_hint(&self) -> Option<&str>

Architecture name from the underlying file (general.architecture for GGUF, None for safetensors). Drain-style consumers use this to pick an arch-specific reverse name mapping when the canonical HF name depends on the model family (e.g. Gemma 2’s 4 norms per layer don’t share the Llama 2-norm reverse alias).