Struct WeightMap 

pub struct WeightMap { /* private fields */ }

Map of tensor name → (f32 data, shape).

Implementations

impl WeightMap

pub fn from_weight_loader(loader: &mut dyn WeightLoader) -> Result<Self>

Drain every tensor from any WeightLoader (safetensors or GGUF).

pub fn from_weight_loader_dequant_all( loader: &mut dyn WeightLoader, ) -> Result<Self>

Force-dequantize every tensor (including K-quants) into F32 and drop it in the map. Use when a family runner doesn’t have a packed-matmul lowering yet but still wants to load GGUFs whose trunk weights are K-quant. Trades memory (4× larger than the packed bytes) for correctness — every tensor goes through WeightLoader::take(...) which dequantizes on the fly.

pub fn drain_loader( loader: &mut dyn WeightLoader, policy: WeightDrainPolicy, ) -> Result<(Self, Vec<NamedPackedWeight>)>

Drain with policy; returns packed K-quants separately when the loader supports take_packed.

pub fn from_resolved_path(path: &Path) -> Result<Self>

Resolve a file or weights directory, then load (safetensors or GGUF).

pub fn from_resolved_safetensors_only(path: &Path, runner: &str) -> Result<Self>

Resolve path; reject .gguf with a hint naming the right runner.

pub fn from_file(path: &str) -> Result<Self>

Load weights from a safetensors file. Auto-converts bf16/f16 to f32.

pub fn from_file_excluding( path: &str, exclude: &HashSet<String>, ) -> Result<Self>

Load weights, skipping tensor names present in exclude (saves RAM when bf16/NVFP4 linears are loaded separately for GPU upload).

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

Take a tensor by name (removes from map). Returns (data, shape).

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

Take and transpose a 2D weight: [out, in] → [in, out] for row-major matmul.

pub fn has(&self, key: &str) -> bool

Check if a key exists.

pub fn keys(&self) -> impl Iterator<Item = &str>

List all keys.

pub fn len(&self) -> usize

Number of tensors remaining.

pub fn is_empty(&self) -> bool

pub fn from_tensors(tensors: HashMap<String, (Vec<f32>, Vec<usize>)>) -> Self

Create from pre-built HashMap (for testing without safetensors files).

pub fn snapshot_from_path(path: &str) -> Result<F32WeightSnapshot>

Drain all tensors into a snapshot map (for runners that rebuild graphs per shape).

pub fn from_safetensors_dir_selected( dir: &Path, want: &HashSet<String>, ) -> Result<Self>

Load only tensors whose names appear in want (HF sharded checkpoints).

pub fn from_safetensors_dir(dir: &Path) -> Result<Self>

Load and merge every *.safetensors file in dir (e.g. HF text_encoder/).

pub fn remap_keys<F>(&mut self, f: F)

where F: FnMut(String) -> String,

Rename keys in-place (e.g. strip model. HuggingFace prefix).

pub fn get(&self, key: &str) -> Option<(&[f32], &[usize])>

Borrow tensor data + shape without removing from the map.

pub fn merge_add_weight(&mut self, key: &str, delta: &[f32]) -> Result<()>

Element-wise add delta into an existing rank-2 weight (PyTorch [out, in] layout).

Trait Implementations

impl WeightLoader for WeightMap

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

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

fn len(&self) -> usize

Number of distinct weights in the file.

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

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

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

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.

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

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

fn is_empty(&self) -> bool

fn take_packed(&mut self, key: &str) -> Result<Option<PackedWeightTensor>>

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

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

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

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