cherenkov_model_data/
mlx.rs

1//! MLX affine conventions over a safetensors inventory. This adapter does not
2//! interpret architecture-specific names or choose execution precision.
3use crate::*;
4use anyhow::{Context, Result, ensure};
5use serde_json::{Map, Value};
6use std::collections::HashMap;
7
8/// Combine MLX weight, scale, and bias entries into logical affine tensors.
9/// Per-layer quantization settings override global settings. Unknown encodings
10/// retain their metadata and data spans as opaque tensors; no weights are decoded.
11pub fn tensors(inventory: &Inventory, config: &Value) -> Result<Vec<Tensor>> {
12    let entries: HashMap<_, _> = inventory
13        .tensors
14        .iter()
15        .map(|t| (t.name.as_str(), t))
16        .collect();
17    let mut tensors = Vec::new();
18
19    for tensor in &inventory.tensors {
20        if is_component(&entries, &tensor.name) {
21            continue;
22        }
23
24        let Some(prefix) = tensor.name.strip_suffix(".weight") else {
25            tensors.push(tensor.clone());
26
27            continue;
28        };
29        let TensorEncoding::Dense { tensor: codes } = &tensor.encoding else {
30            tensors.push(tensor.clone());
31
32            continue;
33        };
34
35        if codes.dtype != Dtype::U32 {
36            tensors.push(tensor.clone());
37
38            continue;
39        }
40
41        let get = |suffix| -> Result<StoredTensor> {
42            let entry = entries
43                .get(format!("{prefix}.{suffix}").as_str())
44                .with_context(|| format!("{prefix}: missing {suffix}"))?;
45            let TensorEncoding::Dense { tensor } = &entry.encoding else {
46                anyhow::bail!("{prefix}: unsupported {suffix} storage");
47            };
48
49            Ok(tensor.clone())
50        };
51        let quantization = quantization(config, prefix);
52        let bits = quantization["bits"].as_u64();
53        let mode = quantization["mode"].as_str().unwrap_or("affine");
54
55        if !matches!(bits, Some(2 | 4 | 8)) || mode != "affine" {
56            let mut data = vec![codes.data.clone()];
57
58            for suffix in ["scales", "biases"] {
59                if let Some(entry) = entries.get(format!("{prefix}.{suffix}").as_str()) {
60                    data.extend(entry.encoding.data().into_iter().cloned());
61                }
62            }
63
64            tensors.push(Tensor::new(
65                tensor.name.clone(),
66                tensor.role,
67                None,
68                TensorEncoding::Opaque {
69                    name: "unrecognized MLX weight encoding".into(),
70                    metadata: quantization,
71                    data,
72                },
73            ));
74
75            continue;
76        }
77
78        let scales = get("scales")?;
79        let biases = get("biases")?;
80        let bits = bits.context("missing quantization bits")? as u8;
81
82        tensors.push(affine(tensor, codes.clone(), scales, biases, bits)?);
83    }
84
85    Ok(tensors)
86}
87
88fn affine(
89    source: &Tensor,
90    codes: StoredTensor,
91    scales: StoredTensor,
92    biases: StoredTensor,
93    bits: u8,
94) -> Result<Tensor> {
95    let mut shape = codes.shape.clone();
96    let axis = shape
97        .len()
98        .checked_sub(1)
99        .context("scalar quantized weight")?;
100    shape[axis] = shape[axis]
101        .checked_mul(32 / u64::from(bits))
102        .context("quantized shape overflow")?;
103
104    ensure!(
105        scales.shape.len() == shape.len(),
106        "affine scale rank mismatch"
107    );
108
109    let groups = scales.shape[axis];
110
111    ensure!(
112        groups > 0 && shape[axis].is_multiple_of(groups),
113        "invalid affine scale count"
114    );
115
116    let encoding = TensorEncoding::Affine {
117        bits,
118        group_size: shape[axis] / groups,
119        group_axis: axis,
120        packing: BitPacking::LowFirstU32,
121        codes,
122        scales,
123        offset: AffineOffset::Bias { tensor: biases },
124    };
125    let tensor = Tensor::new(source.name.clone(), source.role, Some(shape), encoding);
126
127    tensor.validate()?;
128
129    Ok(tensor)
130}
131
132fn is_component(entries: &HashMap<&str, &Tensor>, name: &str) -> bool {
133    let Some(prefix) = name
134        .strip_suffix(".scales")
135        .or_else(|| name.strip_suffix(".biases"))
136    else {
137        return false;
138    };
139
140    entries
141        .get(format!("{prefix}.weight").as_str())
142        .is_some_and(|entry| {
143            matches!(
144                &entry.encoding, TensorEncoding::Dense { tensor } if tensor.dtype == Dtype::U32
145            )
146        })
147}
148
149fn quantization(config: &Value, prefix: &str) -> Value {
150    let Some(root) = config
151        .get("quantization")
152        .or_else(|| config.get("quantization_config"))
153    else {
154        return Value::Null;
155    };
156    let mut fields = Map::new();
157
158    for name in ["bits", "group_size", "mode"] {
159        if let Some(value) = root.get(name) {
160            fields.insert(name.into(), value.clone());
161        }
162    }
163
164    if let Some(local) = root.get(prefix).and_then(Value::as_object) {
165        fields.extend(local.clone());
166    }
167
168    Value::Object(fields)
169}
cherenkov_model_data/mlx.rs

cherenkov_model_data/
mlx.rs
