rlx_pkg/
weights_index.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! `weights/index.json` tensor catalog.
6
7use crate::tier::{Codec, StorageTier};
8use serde::{Deserialize, Serialize};
9
10/// One weight tensor in the package catalog.
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
12pub struct WeightEntry {
13    pub name: String,
14    pub shape: Vec<usize>,
15    /// `f32` | `gguf_q8_0` | `gguf_tq2_0` | …
16    pub scheme: String,
17    #[serde(default = "default_layout")]
18    pub layout: String,
19    /// Path relative to package root (`__flat__` for flat packs).
20    pub shard: String,
21    pub offset: u64,
22    /// Stored (on-disk) byte length.
23    pub length: u64,
24    /// Uncompressed length when compressed; defaults to `length`.
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub raw_length: Option<u64>,
27    #[serde(default)]
28    pub tier: StorageTier,
29    #[serde(default)]
30    pub codec: Codec,
31    #[serde(default, skip_serializing_if = "Option::is_none")]
32    pub rank: Option<u32>,
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub checksum: Option<String>,
35}
36
37fn default_layout() -> String {
38    "row_major".into()
39}
40
41impl WeightEntry {
42    pub fn is_mmapable(&self) -> bool {
43        self.tier == StorageTier::Hot && self.codec == Codec::None
44    }
45}
46
47/// Weight catalog file.
48#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
49pub struct WeightsIndex {
50    pub tensors: Vec<WeightEntry>,
51}
52
53impl WeightsIndex {
54    pub fn get(&self, name: &str) -> Option<&WeightEntry> {
55        self.tensors.iter().find(|t| t.name == name)
56    }
57
58    pub fn names(&self) -> impl Iterator<Item = &str> {
59        self.tensors.iter().map(|t| t.name.as_str())
60    }
61}
rlx_pkg/weights_index.rs

rlx_pkg/
weights_index.rs
