zyx/

module.rs

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only WITH Classpath-exception-2.0

use std::{collections::HashMap, ffi::OsStr, fs::File, io::Seek, path::Path};

use crate::{DType, Map, Tensor, ZyxError, shape::Dim};

/// Module trait
pub trait Module {
    /// Iterate over all tensors immutably
    fn iter(&self) -> impl Iterator<Item = &Tensor>;

    /// Iterate over all tensors mutably
    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Tensor>;

    /// Iterate over tensors without consuming the module
    fn iter_tensors(&self) -> impl Iterator<Item = (String, &Tensor)>;

    /// From tensors
    fn iter_tensors_mut(&mut self) -> impl Iterator<Item = (String, &mut Tensor)>;

    /// Set parameters, removes them from params, skips parameters that are not found in params.
    fn set_params(&mut self, params: &mut HashMap<String, Tensor>) {
        for (label, tensor) in self.iter_tensors_mut() {
            if let Some(param) = params.remove(&label) {
                *tensor = param;
            }
        }
    }

    /// Save tensors or modules to a file determined by file extension.
    /// Currently only safetensors is supported format.
    ///
    /// # Errors
    ///
    /// Errors if tensors failed to realize or failed to save to disk.
    fn save(&self, path: impl AsRef<Path>) -> Result<(), ZyxError> {
        use std::fmt::Write;
        use std::io::Write as IOWrite;
        let mut f = File::create(path)?;
        let mut header = String::from("{");
        let mut begin = 0;
        for (label, tensor) in self.iter_tensors() {
            let dtype = tensor.dtype();
            write!(header, "\"{label}\":{{").unwrap();
            write!(header, "\"dtype\":\"{}\",", dtype.safetensors()).unwrap();
            let mut st_shape = format!("{:?}", tensor.resolve_shape());
            st_shape.retain(|c| !c.is_whitespace());
            write!(header, "\"shape\":{st_shape},").unwrap();
            let size = tensor.numel().item::<Dim>() * Dim::from(dtype.bit_size() / 8);
            write!(header, "\"data_offsets\":[{},{}]", begin, begin + size).unwrap();
            begin += size;
            write!(header, "}},").unwrap();
        }
        header.pop();
        write!(header, "}}").unwrap();
        let header_bytes = header.as_bytes();
        f.write_all(&(header_bytes.len() as i64).to_le_bytes())?;
        f.write_all(header_bytes)?;
        for tensor in self.iter() {
            f.write_all(&tensor.to_le_bytes()?)?;
        }
        Ok(())
    }

    /// Save a single tensor to a `.npy` file (numpy array format).
    /// Mirrors `load_numpy`: little-endian, C order (Fortran order
    /// is never written). Header is padded so data starts at a 64-byte
    /// boundary, like numpy >= 1.9. Numpy files hold a single array, so
    /// saving a module with more than one tensor is an error.
    ///
    /// # Errors
    ///
    /// Errors if the module holds more than one tensor, if the tensor
    /// failed to realize or failed to save to disk.
    fn save_numpy(&self, path: impl AsRef<Path>) -> Result<(), ZyxError> {
        use std::io::Write as IOWrite;
        let mut tensors = self.iter_tensors();
        let (label, tensor) = match (tensors.next(), tensors.next()) {
            (Some((label, tensor)), None) => (label, tensor),
            (None, _) => return Err(ZyxError::parse_error("Cannot save empty module to numpy: no tensors.".into())),
            (Some((l0, _)), Some((l1, _))) => {
                return Err(ZyxError::parse_error(
                    format!(
                        "Cannot save module to numpy: numpy files hold a single array, module has tensors '{l0}' and '{l1}' (and possibly more)."
                    )
                    .into(),
                ));
            }
        };
        let _ = label;
        let descr = match tensor.dtype() {
            DType::F32 => "<f4",
            DType::F64 => "<f8",
            DType::F16 => "<f2",
            DType::I8 => "|i1",
            DType::I16 => "<i2",
            DType::I32 => "<i4",
            DType::I64 => "<i8",
            DType::U8 => "|u1",
            DType::U16 => "<u2",
            DType::BF16 => todo!("BF16 has no numpy dtype"),
            DType::U32 => todo!("u4 numpy arrays"),
            DType::U64 => todo!("u8 numpy arrays"),
            DType::Bool => todo!("Bool numpy arrays"),
            DType::F8E4M3 => todo!("F8E4M3 numpy arrays"),
            DType::F8E5M2 => todo!("F8E5M2 numpy arrays"),
        };
        let dims = tensor.resolve_shape();
        let shape_str = format!("({})", dims.iter().map(|d| d.to_string()).collect::<Vec<_>>().join(", "));
        let mut header = format!("{{'descr': '{descr}', 'fortran_order': False, 'shape': {shape_str}, }}");
        // magic(6) + version(2) + header_len(2) + header + '\n' must be a
        // multiple of 64.
        let total = 6 + 2 + 2 + header.len() + 1;
        header.extend(core::iter::repeat(' ').take((64 - total % 64) % 64));
        header.push('\n');
        let mut f = File::create(path)?;
        f.write_all(b"\x93NUMPY")?;
        f.write_all(&[1u8, 0u8])?;
        f.write_all(&(header.len() as u16).to_le_bytes())?;
        f.write_all(header.as_bytes())?;
        f.write_all(&tensor.to_le_bytes()?)?;
        Ok(())
    }
}

/// GGUF metadata value.
///
/// Maps one-to-one onto the GGUF file format metadata types
/// (`TYPE_INT8` .. `TYPE_UINT64`, `TYPE_F32`, `TYPE_F64`, `TYPE_BOOL`,
/// `TYPE_STRING`), plus `GGUF_ARRAY`, whose elements are recursively
/// `GGUFMetadataValue`s.
#[allow(unused)]
pub enum GGUFMetadataValue {
    /// Unsigned 8-bit integer
    Uint8(u8),
    /// Signed 8-bit integer
    Int8(i8),
    /// Unsigned 16-bit integer
    Uint16(u16),
    /// Signed 16-bit integer
    Int16(i16),
    /// Unsigned 32-bit integer
    Uint32(u32),
    /// Signed 32-bit integer
    Int32(i32),
    /// Unsigned 64-bit integer
    Uint64(u64),
    /// Signed 64-bit integer
    Int64(i64),
    /// 32-bit floating-point number
    Float32(f32),
    /// 64-bit floating-point number
    Float64(f64),
    /// Boolean value
    Bool(bool),
    /// UTF-8 string
    String(String),
    /// Array of arbitrary metadata values
    Array(Box<[GGUFMetadataValue]>),
}

impl<S: std::hash::BuildHasher + Default> Module for HashMap<String, Tensor, S> {
    fn iter(&self) -> impl Iterator<Item = &Tensor> {
        self.values()
    }

    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Tensor> {
        self.values_mut()
    }

    fn iter_tensors(&self) -> impl Iterator<Item = (String, &Tensor)> {
        self.iter().map(|(k, v): (&String, &Tensor)| (k.clone(), v))
    }

    fn iter_tensors_mut(&mut self) -> impl Iterator<Item = (String, &mut Tensor)> {
        self.iter_mut().map(|(k, v): (&String, &mut Tensor)| (k.clone(), v))
    }
}

impl Module for Vec<Tensor> {
    #[allow(clippy::into_iter_on_ref)] // into_iter on &Vec/&mut Vec is the existing pattern; changing resolution risks recursion
    fn iter(&self) -> impl Iterator<Item = &Tensor> {
        self.into_iter()
    }

    #[allow(clippy::into_iter_on_ref)] // into_iter on &Vec/&mut Vec is the existing pattern; changing resolution risks recursion
    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Tensor> {
        self.into_iter()
    }

    fn iter_tensors(&self) -> impl Iterator<Item = (String, &Tensor)> {
        self.iter().map(|t: &Tensor| (format!("{}", t.id()), t))
    }

    fn iter_tensors_mut(&mut self) -> impl Iterator<Item = (String, &mut Tensor)> {
        self.iter_mut().map(|t: &mut Tensor| (format!("{}", t.id()), t))
    }
}

impl<M0: Module, M1: Module> Module for (M0, M1) {
    fn iter(&self) -> impl Iterator<Item = &Tensor> {
        self.0.iter().chain(self.1.iter())
    }

    fn iter_mut(&mut self) -> impl Iterator<Item = &mut Tensor> {
        self.0.iter_mut().chain(self.1.iter_mut())
    }

    fn iter_tensors(&self) -> impl Iterator<Item = (String, &Tensor)> {
        self.0.iter_tensors().chain(self.1.iter_tensors())
    }

    fn iter_tensors_mut(&mut self) -> impl Iterator<Item = (String, &mut Tensor)> {
        self.0.iter_tensors_mut().chain(self.1.iter_tensors_mut())
    }
}

impl Tensor {
    /// Load module from path. This function will determine the filetype based on file extension.
    ///
    /// # Errors
    ///
    /// Errors if the path has no or an unknown extension, if loading from disk failed
    /// or if loaded tensors could not be allocated to device.
    pub fn load(path: impl AsRef<Path>) -> Result<HashMap<String, Tensor>, ZyxError>
    where
        Self: Sized,
    {
        let e = path.as_ref().extension().and_then(OsStr::to_str);
        match e {
            Some("safetensors") => Self::load_safetensors(path),
            Some("gguf") => Ok(Self::load_gguf(path)?.1),
            Some(other) => Err(ZyxError::parse_error(
                format!("Unknown file extension '{other}'. Zyx currently supports only safetensors and gguf formats.").into(),
            )),
            None => Err(ZyxError::parse_error(
                format!("Cannot determine file type: '{}' has no extension. Zyx currently supports only safetensors and gguf formats.", path.as_ref().display()).into(),
            )),
        }
    }

    /// Load gguf module from path
    /// First returned value is metadata, second returned value are named tensors
    /// # Errors
    /// read failure
    #[allow(clippy::missing_panics_doc)]
    #[allow(clippy::type_complexity)]
    pub fn load_gguf(path: impl AsRef<Path>) -> Result<(HashMap<String, GGUFMetadataValue>, HashMap<String, Tensor>), ZyxError> {
        use std::io::Read;
        let mut f = std::fs::File::open(&path)?;
        let mut magic = [0; 4];
        f.read_exact(&mut magic)?;
        if magic != *b"GGUF" {
            if magic == *b"FUGG" {
                return Err(ZyxError::parse_error(
                    "GGUF data seems to be stored in big endian order. Only little endian is supported for GGUF in zyx.".into(),
                ));
            }
            return Err(ZyxError::parse_error(format!("Unknown GGUF magic: {magic:?}. Please check your file.").into()));
        }
        let mut version_bytes = [0; 4];
        f.read_exact(&mut version_bytes)?;
        let version = u32::from_le_bytes(version_bytes);
        //println!("File size is {} bytes", f.metadata()?.len());
        let mut tensor_count = [0u8; 8];
        f.read_exact(&mut tensor_count)?;
        let tensor_count = u64::from_le_bytes(tensor_count);
        let mut metadata_kv_count = [0u8; 8];
        f.read_exact(&mut metadata_kv_count)?;
        let metadata_kv_count = usize::try_from(u64::from_le_bytes(metadata_kv_count))
            .map_err(|e| ZyxError::parse_error(format!("Failed to parse tensor count in GGUF file. {e}").into()))?;

        let mut metadata = HashMap::new();
        for _ in 0..metadata_kv_count {
            // First string key, (len u64, chars),
            let mut metadata_key_len = [0; 8];
            f.read_exact(&mut metadata_key_len)?;
            let metadata_key_len = u64::from_le_bytes(metadata_key_len);
            let mut metadata_key_bytes = vec![0u8; usize::try_from(metadata_key_len).unwrap()];
            f.read_exact(&mut metadata_key_bytes)?;
            let metadata_key = String::from_utf8(metadata_key_bytes)
                .map_err(|e| ZyxError::parse_error(format!("GGUF metadata key is not valid UTF-8: {e}").into()))?;

            // Then metadata value type (u32 in GGUF v3, u8 in v1/v2).
            // Then we the value itself.
            let metadata_value_type = if version >= 3 {
                let mut buf = [0; 4];
                f.read_exact(&mut buf)?;
                u32::from_le_bytes(buf)
            } else {
                let mut buf = [0; 1];
                f.read_exact(&mut buf)?;
                u32::from(u8::from_le_bytes(buf))
            };
            let metadata_value = match metadata_value_type {
                0 => {
                    let mut buf = [0; 1];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Uint8(u8::from_le_bytes(buf))
                }
                1 => {
                    let mut buf = [0; 1];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Int8(i8::from_le_bytes(buf))
                }
                2 => {
                    let mut buf = [0; 2];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Uint16(u16::from_le_bytes(buf))
                }
                3 => {
                    let mut buf = [0; 2];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Int16(i16::from_le_bytes(buf))
                }
                4 => {
                    let mut buf = [0; 4];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Uint32(u32::from_le_bytes(buf))
                }
                5 => {
                    let mut buf = [0; 4];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Int32(i32::from_le_bytes(buf))
                }
                6 => {
                    let mut buf = [0; 4];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Float32(f32::from_le_bytes(buf))
                }
                7 => {
                    let mut buf = [0; 1];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Bool(buf[0] != 0)
                }
                8 => {
                    let mut str_len = [0; 8];
                    f.read_exact(&mut str_len)?;
                    let str_len = u64::from_le_bytes(str_len);
                    let mut s_bytes = vec![0u8; usize::try_from(str_len).unwrap()];
                    f.read_exact(&mut s_bytes)?;
                    let s = String::from_utf8(s_bytes)
                        .map_err(|e| ZyxError::parse_error(format!("GGUF metadata string is not valid UTF-8: {e}").into()))?;
                    GGUFMetadataValue::String(s)
                }
                9 => {
                    let mut arr_type_buf = [0; 4];
                    f.read_exact(&mut arr_type_buf)?;
                    let elem_type = u32::from_le_bytes(arr_type_buf);
                    let mut arr_len_buf = [0; 8];
                    f.read_exact(&mut arr_len_buf)?;
                    let arr_len = u64::from_le_bytes(arr_len_buf);
                    let mut items = Vec::with_capacity(usize::try_from(arr_len).unwrap());
                    for _ in 0..arr_len {
                        let item = match elem_type {
                            0 => {
                                let mut buf = [0; 1];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Uint8(u8::from_le_bytes(buf))
                            }
                            1 => {
                                let mut buf = [0; 1];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Int8(i8::from_le_bytes(buf))
                            }
                            2 => {
                                let mut buf = [0; 2];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Uint16(u16::from_le_bytes(buf))
                            }
                            3 => {
                                let mut buf = [0; 2];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Int16(i16::from_le_bytes(buf))
                            }
                            4 => {
                                let mut buf = [0; 4];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Uint32(u32::from_le_bytes(buf))
                            }
                            5 => {
                                let mut buf = [0; 4];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Int32(i32::from_le_bytes(buf))
                            }
                            6 => {
                                let mut buf = [0; 4];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Float32(f32::from_le_bytes(buf))
                            }
                            7 => {
                                let mut buf = [0; 1];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Bool(buf[0] != 0)
                            }
                            8 => {
                                let mut item_len = [0; 8];
                                f.read_exact(&mut item_len)?;
                                let item_len = u64::from_le_bytes(item_len);
                                let mut item_bytes = vec![0u8; usize::try_from(item_len).unwrap()];
                                f.read_exact(&mut item_bytes)?;
                                let item = String::from_utf8(item_bytes).map_err(|e| {
                                    ZyxError::parse_error(format!("GGUF array element string is not valid UTF-8: {e}").into())
                                })?;
                                GGUFMetadataValue::String(item)
                            }
                            10 => {
                                let mut buf = [0; 8];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Uint64(u64::from_le_bytes(buf))
                            }
                            11 => {
                                let mut buf = [0; 8];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Int64(i64::from_le_bytes(buf))
                            }
                            12 => {
                                let mut buf = [0; 8];
                                f.read_exact(&mut buf)?;
                                GGUFMetadataValue::Float64(f64::from_le_bytes(buf))
                            }
                            x => todo!("GGUF array element type {x} not supported"),
                        };
                        items.push(item);
                    }
                    GGUFMetadataValue::Array(items.into_boxed_slice())
                }
                10 => {
                    let mut buf = [0; 8];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Uint64(u64::from_le_bytes(buf))
                }
                11 => {
                    let mut buf = [0; 8];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Int64(i64::from_le_bytes(buf))
                }
                12 => {
                    let mut buf = [0; 8];
                    f.read_exact(&mut buf)?;
                    GGUFMetadataValue::Float64(f64::from_le_bytes(buf))
                }
                x => todo!("GGUF metadata type {x} not supported"),
            };
            metadata.insert(metadata_key, metadata_value);
        }

        // First we read the whole description of tensors
        let mut tensor_header = Map::default();
        for _ in 0..tensor_count {
            // name
            let mut tensor_name_len = [0; 8];
            f.read_exact(&mut tensor_name_len)?;
            let tensor_name_len = u64::from_le_bytes(tensor_name_len);
            let mut tensor_name_bytes = vec![0u8; usize::try_from(tensor_name_len).unwrap()];
            f.read_exact(&mut tensor_name_bytes)?;
            let tensor_name = String::from_utf8(tensor_name_bytes)
                .map_err(|e| ZyxError::parse_error(format!("GGUF tensor name is not valid UTF-8: {e}").into()))?;

            // rank (number of dimensions)
            let mut rank = [0; 4];
            f.read_exact(&mut rank)?;
            let rank = u32::from_le_bytes(rank);

            // shape (NOTE there is no explicit check for endiannes here)
            let mut shape = vec![0u8; rank as usize * 8];
            f.read_exact(&mut shape)?;
            let shape: Vec<Dim> =
                shape.chunks_exact(8).map(|x| i64::from_le_bytes([x[0], x[1], x[2], x[3], x[4], x[5], x[6], x[7]])).collect();

            // dtype
            let mut dtype = [0; 4];
            f.read_exact(&mut dtype)?;
            let dtype = u32::from_le_bytes(dtype);
            // Q4_K (gguf type 12) loads as raw super-blocks: [num_blocks, 144]
            // U8. Each 144B block holds 256 weights (d, dmin, 12 scale bytes,
            // 128B of nibbles, llama.cpp `block_q4_K`). Element dims are not
            // representable at sub-byte granularity, so the caller derives
            // (rows, cols) from the model hyperparams + tensor name.
            let (dtype, shape) = match dtype {
                0 => (DType::F32, shape),
                1 => (DType::F16, shape),
                24 => (DType::I8, shape),
                25 => (DType::I16, shape),
                26 => (DType::I32, shape),
                27 => (DType::I64, shape),
                28 => (DType::F64, shape),
                12 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "Q4_K tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 144])
                }
                // Q8_0: block_q8_0, 34B per 32 (static_assert: half + QK8_0).
                8 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 32 == 0, "Q8_0 tensor {tensor_name} has {numel} elements, not a multiple of 32");
                    (DType::U8, vec![numel / 32, 34])
                }
                // Q3_K: block_q3_K, 110B per 256 (half + 64 + 32 + 12).
                11 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "Q3_K tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 110])
                }
                // Q5_K: block_q5_K, 176B per 256 (2*half + 12 + 128 + 32).
                13 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "Q5_K tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 176])
                }
                // Q6_K: block_q6_K, 210B per 256 (half + 16 + 192).
                14 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "Q6_K tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 210])
                }
                // IQ4_NL: block_iq4_nl, 18B per 32 (half + QK4_NL/2).
                20 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 32 == 0, "IQ4_NL tensor {tensor_name} has {numel} elements, not a multiple of 32");
                    (DType::U8, vec![numel / 32, 18])
                }
                // IQ3_S: block_iq3_s, 110B per 256 (half + 104 + 4).
                21 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "IQ3_S tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 110])
                }
                // IQ4_XS: block_iq4_xs, 136B per 256 (half + u16 + 4 + 128).
                23 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "IQ4_XS tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 136])
                }
                // Q4_0 (gguf type 2) loads as raw super-blocks: [num_blocks, 18]
                // U8. Each 18B block holds 32 weights (half d + 16B of nibbles,
                // llama.cpp `block_q4_0`).
                2 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 32 == 0, "Q4_0 tensor {tensor_name} has {numel} elements, not a multiple of 32");
                    (DType::U8, vec![numel / 32, 18])
                }
                // Q4_1 (gguf type 3): block_q4_1, 20B per 32 (2*half + 16B nibbles).
                3 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 32 == 0, "Q4_1 tensor {tensor_name} has {numel} elements, not a multiple of 32");
                    (DType::U8, vec![numel / 32, 20])
                }
                // Q5_0 (gguf type 6): block_q5_0, 22B per 32 (half + 4B qh + 16B qs).
                6 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 32 == 0, "Q5_0 tensor {tensor_name} has {numel} elements, not a multiple of 32");
                    (DType::U8, vec![numel / 32, 22])
                }
                // Q5_1 (gguf type 7): block_q5_1, 24B per 32 (2*half + 4B qh + 16B qs).
                7 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 32 == 0, "Q5_1 tensor {tensor_name} has {numel} elements, not a multiple of 32");
                    (DType::U8, vec![numel / 32, 24])
                }
                // Q8_1 (gguf type 9): block_q8_1, 36B per 32 (2*half + 32B qs).
                9 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 32 == 0, "Q8_1 tensor {tensor_name} has {numel} elements, not a multiple of 32");
                    (DType::U8, vec![numel / 32, 36])
                }
                // Q2_K (gguf type 10): block_q2_K, 84B per 256 (2*half + 16B scales + 64B qs).
                10 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "Q2_K tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 84])
                }
                // Q8_K (gguf type 15): block_q8_K, 292B per 256 (float + 256B qs + 16 i16 sums).
                15 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "Q8_K tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 292])
                }
                // IQ2_XXS (gguf type 16): block_iq2_xxs, 66B per 256 (half + 32 u16 qs).
                16 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "IQ2_XXS tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 66])
                }
                // IQ2_XS (gguf type 17): block_iq2_xs, 74B per 256 (half + 64B qs + 8B scales).
                17 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "IQ2_XS tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 74])
                }
                // IQ3_XXS (gguf type 18): block_iq3_xxs, 98B per 256 (half + 96B qs).
                18 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "IQ3_XXS tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 98])
                }
                // IQ1_S (gguf type 19): block_iq1_s, 50B per 256 (half + 32B qs + 16B qh).
                19 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "IQ1_S tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 50])
                }
                // IQ2_S (gguf type 22): block_iq2_s, 82B per 256 (half + 64B qs + 8B qh + 8B scales).
                22 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "IQ2_S tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 82])
                }
                // IQ1_M (gguf type 29): block_iq1_m, 56B per 256 (32B qs + 16B qh + 8B scales, no fp scale).
                29 => {
                    let numel: Dim = shape.iter().product();
                    debug_assert!(numel % 256 == 0, "IQ1_M tensor {tensor_name} has {numel} elements, not a multiple of 256");
                    (DType::U8, vec![numel / 256, 56])
                }
                x => todo!("GGUF dtype {x} is not supported by zyx yet."),
            };

            // offset (position in file)
            let mut offset = [0; 8];
            f.read_exact(&mut offset)?;
            let offset = u64::from_le_bytes(offset);

            tensor_header.insert(tensor_name, (shape, dtype, offset));
        }

        // GGUF tensor offsets are relative to the data section, which starts
        // right after the tensor infos, aligned up to `general.alignment`
        // (spec default 32). The offsets must not be used as raw file
        // offsets.
        let alignment = match metadata.get("general.alignment") {
            Some(GGUFMetadataValue::Uint32(a)) => (*a as usize).max(1),
            Some(_) => todo!("general.alignment must be Uint32"),
            None => 32,
        };
        let data_start = f.stream_position()? as usize;
        let data_start = data_start.div_ceil(alignment) * alignment;

        let mut progress_bar = if crate::debug_mask().dev() {
            println!("Loading tensors from safetensors file");
            let bar = crate::progress::ProgressBar::new(tensor_count);
            Some(bar)
        } else {
            None
        };

        let mut tensors = HashMap::new();
        for (name, (shape, dtype, offset)) in tensor_header {
            if let Some(progress_bar) = &mut progress_bar {
                progress_bar.inc(1, &format!("{name}, {shape:?}, {dtype}"));
            }
            tensors.insert(name, Tensor::from_path(shape, dtype, &path, (data_start as u64) + offset)?);
        }
        Ok((metadata, tensors))
    }

    /// Load a single `.npy` array from path.
    ///
    /// Reads the array lazily from disk (no host copy until realize), like
    /// [`Self::load_gguf`] and [`Self::load_safetensors`]. Supports little
    /// -endian numeric dtypes; big-endian files, Fortran order and non
    /// -numeric dtypes are loud errors, never guesses.
    ///
    /// # Errors
    /// Errors if the path does not exist, IO failed, or the file uses an
    /// unsupported dtype, byte order or memory order.
    pub fn load_numpy(path: impl AsRef<Path>) -> Result<Tensor, ZyxError> {
        use std::io::Read;
        let path = path.as_ref();
        let mut f = File::open(path)?;
        let mut magic = [0; 6];
        f.read_exact(&mut magic)?;
        if magic != *b"\x93NUMPY" {
            return Err(ZyxError::parse_error(format!("Unknown numpy magic: {magic:?} in {path:?}").into()));
        }
        let mut ver = [0; 2];
        f.read_exact(&mut ver)?;
        // v1.0 header len is u16, v2.0+ is u32.
        let header_len = match ver[0] {
            1 => {
                let mut buf = [0; 2];
                f.read_exact(&mut buf)?;
                u16::from_le_bytes(buf) as usize
            }
            2 | 3 => {
                let mut buf = [0; 4];
                f.read_exact(&mut buf)?;
                u32::from_le_bytes(buf) as usize
            }
            x => return Err(ZyxError::parse_error(format!("Unsupported numpy version {x} in {path:?}").into())),
        };
        let mut header = vec![0u8; header_len];
        f.read_exact(&mut header)?;
        let header = String::from_utf8(header)
            .map_err(|e| ZyxError::parse_error(format!("numpy header is not valid UTF-8: {e} in {path:?}").into()))?;
        // Header is a python dict literal: {'descr': '<f4', 'fortran_order': False, 'shape': (2, 3), }
        let field = |key: &str| -> Option<String> {
            let start = header.find(&format!("'{key}':"))? + key.len() + 4;
            // Value ends at the next top-level ',' or '}'.
            let rest = &header[start..];
            let end = rest.find(|c| c == ',' || c == '}').unwrap_or(rest.len());
            Some(rest[..end].trim().to_string())
        };
        let descr =
            field("descr").ok_or_else(|| ZyxError::parse_error(format!("numpy header missing 'descr' in {path:?}").into()))?;
        let descr = descr.trim_matches(|c| c == '\'' || c == '"').to_string();
        let fortran = field("fortran_order").unwrap_or_default();
        if fortran.contains("True") {
            return Err(ZyxError::parse_error(format!("Fortran-order numpy arrays are not supported: {path:?}").into()));
        }
        // The shape value is a tuple "(2, 3)" containing commas itself, so
        // it cannot go through `field`, which stops at the first ','. Take
        // everything up to the closing '}' of the dict instead.
        let shape_start = header
            .find("'shape':")
            .ok_or_else(|| ZyxError::parse_error(format!("numpy header missing 'shape' in {path:?}").into()))?
            + 8;
        let rest = &header[shape_start..];
        let end = rest.find('}').unwrap_or(rest.len());
        let shape_str = rest[..end].trim().trim_end_matches(',').trim();
        let shape_str = shape_str.trim_matches(|c| c == '(' || c == ')');
        let shape: Vec<Dim> = shape_str
            .split(',')
            .filter(|d| !d.trim().is_empty())
            .map(|d| {
                d.trim()
                    .parse::<Dim>()
                    .map_err(|e| ZyxError::parse_error(format!("Cannot parse numpy shape '{shape_str}': {e} in {path:?}").into()))
            })
            .collect::<Result<_, ZyxError>>()?;
        let dtype = match descr.as_str() {
            "<f4" | "|f4" | "f4" => DType::F32,
            "<f2" | "|f2" | "f2" => DType::F16,
            "<f8" | "|f8" | "f8" => DType::F64,
            "<i1" | "|i1" => DType::I8,
            "<i2" | "|i2" => DType::I16,
            "<i4" | "|i4" => DType::I32,
            "<i8" | "|i8" => DType::I64,
            "|u1" | "<u1" | "u1" => DType::U8,
            "<u2" | "|u2" => DType::U16,
            "<u4" | "|u4" => todo!("u4 numpy arrays"),
            "<u8" | "|u8" => todo!("u8 numpy arrays"),
            x => todo!("numpy dtype '{x}' is not supported ({path:?})"),
        };
        // numpy >= 1.9 pads the header so data starts at a 64-byte boundary;
        // the padding is already counted in header_len, so the stream
        // position after the header is the data start.
        let data_start = f.stream_position()?;
        Tensor::from_path(shape, dtype, path, data_start)
    }

    /// Load safetensors module from path
    ///
    /// # Errors
    /// Errors if path does not exist or IO failed for other reasons.
    #[allow(clippy::missing_panics_doc)]
    pub fn load_safetensors(path: impl AsRef<Path>) -> Result<HashMap<String, Tensor>, ZyxError> {
        use std::io::Read;
        let mut f = std::fs::File::open(&path)?;
        //println!("File size is {} bytes", f.metadata()?.len());
        let mut header_len = [0u8; 8];
        f.read_exact(&mut header_len)?;
        let n = usize::try_from(u64::from_le_bytes(header_len))
            .map_err(|e| ZyxError::parse_error(format!("Failed to parse header len in safetensors file. {e}").into()))?;
        let mut header = vec![0u8; n];
        f.read_exact(&mut header)?;
        let header = core::str::from_utf8(&header).map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?;
        let mut text = String::with_capacity(10);
        let mut begin_str = false;
        let mut i = 0;
        let mut tensors = HashMap::default();
        let mut dtype = DType::F32;
        let mut shape = vec![1i64];
        let mut label = String::new();
        let mut metadata = true;
        let mut progress_bar = if crate::debug_mask().dev() {
            println!("Loading tensors from safetensors file");
            let bar = crate::progress::ProgressBar::new(u64::try_from(header.chars().filter(|&c| c == '[').count()).unwrap() / 2);
            Some(bar)
        } else {
            None
        };
        //let mmap = Arc::new(unsafe { memmap2::Mmap::map(&f)? });
        //let mut mptr = mmap.as_ptr();
        //mptr = mptr.wrapping_add(8 + header.len());
        let mut offset = (8 + header.len()) as i64;
        for x in header.chars() {
            // We skip metadata for now
            if metadata && text.starts_with("__metadata__") {
                if x == '}' {
                    text.clear();
                    begin_str = false;
                    metadata = false;
                }
                continue;
            }
            if ['"', '[', ']'].contains(&x) {
                if begin_str {
                    //std::println!("{text}");
                    if i % 7 == 0 {
                        #[allow(clippy::assigning_clones)]
                        {
                            label = text.clone();
                        }
                    } else if i % 7 == 2 {
                        dtype = DType::from_safetensors(&text)?;
                    } else if i % 7 == 4 {
                        shape = text
                            .split(',')
                            .map(|d| {
                                d.parse::<Dim>()
                                    .map_err(|err| ZyxError::parse_error(format!("Cannot parse safetensors shape: {err}").into()))
                            })
                            .collect::<Result<_, ZyxError>>()?;
                    } else if i % 7 == 6 {
                        // TODO assert offsets
                        //println!("Offsets: {text}");
                        let offsets = text
                            .split(',')
                            .map(|offset| {
                                // Whitespace after commas is valid JSON; the
                                // scanner keeps it in `text`, so trim first.
                                offset.trim().parse::<u64>().map_err(|err| {
                                    ZyxError::parse_error(format!("Could not parse safetensors offset: {err}").into())
                                })
                            })
                            .collect::<Result<Vec<_>, ZyxError>>()?;
                        //println!("Offsets: {offsets:?}");
                        let bytes = shape.iter().product::<Dim>() * Dim::from(dtype.bit_size() / 8);
                        if offsets[1] - offsets[0] != bytes as u64 {
                            return Err(ZyxError::parse_error("Safetensors shapes and offsets are incorrect.".into()));
                        }
                        if let Some(bar) = &mut progress_bar {
                            bar.inc(1, &format!("{label}, {shape:?}, {dtype:?}"));
                        }
                        let tensor = Tensor::from_path(shape.clone(), dtype, &path, offset as u64)?;
                        offset += bytes as i64;
                        tensors.insert(label.clone(), tensor);
                    }
                    i += 1;
                    text.clear();
                    begin_str = false;
                } else {
                    text.clear();
                    begin_str = true;
                }
            } else {
                text.push(x);
            }
        }
        Ok(tensors)
    }
}