teeny_quant/
write.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! Writing `.safetensors` output. `teeny-data::safetensors` only supports mmap'd reading, so
18//! this writes directly via the `safetensors` crate rather than extending that (leaf) crate --
19//! see `teenygrad-303.2`.
20
21use std::collections::HashMap;
22use std::path::Path;
23
24use safetensors::Dtype;
25use safetensors::tensor::TensorView;
26
27use crate::error::Result;
28
29/// One tensor to be written: its dtype, shape, and raw little-endian bytes (already in whatever
30/// on-disk representation `dtype` implies -- e.g. packed `U8` nibbles for INT4).
31#[derive(Debug, Clone)]
32pub struct OutputTensor {
33    /// The tensor's on-disk dtype.
34    pub dtype: Dtype,
35    /// The tensor's logical shape (for packed formats like INT4 this is the *packed* shape --
36    /// see `quantization_config` metadata for the true logical shape).
37    pub shape: Vec<usize>,
38    /// Raw little-endian element bytes, `shape.iter().product() * dtype.bitsize() / 8` long.
39    pub data: Vec<u8>,
40}
41
42fn build_views(tensors: &HashMap<String, OutputTensor>) -> Result<HashMap<String, TensorView<'_>>> {
43    let mut views = HashMap::with_capacity(tensors.len());
44    for (name, t) in tensors {
45        views.insert(
46            name.clone(),
47            TensorView::new(t.dtype, t.shape.clone(), &t.data)?,
48        );
49    }
50    Ok(views)
51}
52
53/// Serializes `tensors` plus `metadata` (embedded as the file's string-keyed `__metadata__`
54/// header -- e.g. the `quantization_config` JSON blob, see [`crate::format`]) to an in-memory
55/// `.safetensors` byte buffer.
56pub fn serialize_safetensors(
57    tensors: &HashMap<String, OutputTensor>,
58    metadata: HashMap<String, String>,
59) -> Result<Vec<u8>> {
60    let views = build_views(tensors)?;
61    Ok(safetensors::serialize(views, Some(metadata))?)
62}
63
64/// Serializes `tensors`/`metadata` (see [`serialize_safetensors`]) directly to a `.safetensors`
65/// file at `path`.
66pub fn write_safetensors(
67    path: &Path,
68    tensors: &HashMap<String, OutputTensor>,
69    metadata: HashMap<String, String>,
70) -> Result<()> {
71    let bytes = serialize_safetensors(tensors, metadata)?;
72    std::fs::write(path, bytes)?;
73    Ok(())
74}
teeny_quant/write.rs

teeny_quant/
write.rs
