Implementation:LaurentMazare Tch rs Tensor Npy

Knowledge Sources	LaurentMazare_Tch_rs
Domains	File I/O, NumPy Interoperability, Serialization
Last Updated	2026-02-08 00:00 GMT

Overview

The npy module implements reading and writing of NumPy .npy and .npz file formats, enabling tensor data exchange between Rust and Python/NumPy ecosystems.

Description

This module provides full support for the NumPy file format specification. It handles both single-tensor .npy files and multi-tensor .npz archives (zip files containing .npy entries).

The internal Header struct captures the three fields from the NumPy header: descr (mapped to Kind), fortran_order (boolean), and shape (Vec<i64>). The Header::parse method implements a custom parser for the Python dictionary-like header format (e.g., {'descr': '<f8', 'fortran_order': False, 'shape': (128,), }). The Header::to_string method generates the header string for writing.

Supported data type descriptors include: f2/Half, f4/Float, f8/Double, i4/Int, i8/Int64, i2/Int16, i1/Int8, u1/Uint8, b1/Bool, c8/ComplexFloat, and c16/ComplexDouble. Both NPY format versions 1 and 2 are supported for reading (version 1 uses 2-byte header length, version 2 uses 4-byte).

The read_header function validates the magic string (\x93NUMPY), reads the version bytes, and extracts the header string. Fortran-order arrays are explicitly rejected.

Four public methods are added to Tensor:

• read_npy: Reads a single tensor from a .npy file
• read_npz: Reads named tensors from a .npz archive
• write_npy: Writes a single tensor to a .npy file (version 1 format)
• write_npz: Writes multiple named tensors to a .npz archive using stored (uncompressed) entries

Usage

Use these methods for persisting tensors in a format compatible with Python/NumPy, for data exchange between Rust and Python pipelines, or for loading pre-computed data from NumPy.

Code Reference

Source Location

• Repository: LaurentMazare_Tch_rs
• File: src/tensor/npy.rs

Signature

impl Tensor {
    /// Reads a npy file and returns the stored tensor.
    pub fn read_npy<T: AsRef<Path>>(path: T) -> Result<Tensor, TchError>;

    /// Reads a npz file and returns some named tensors.
    pub fn read_npz<T: AsRef<Path>>(path: T) -> Result<Vec<(String, Tensor)>, TchError>;

    /// Writes a tensor in the npy format.
    pub fn write_npy<T: AsRef<Path>>(&self, path: T) -> Result<(), TchError>;

    /// Writes multiple named tensors to a npz file.
    pub fn write_npz<S: AsRef<str>, T: AsRef<Tensor>, P: AsRef<Path>>(
        ts: &[(S, T)],
        path: P,
    ) -> Result<(), TchError>;
}

Import

use tch::Tensor;

I/O Contract

Method	Input	Output	Error Conditions
read_npy	File path to .npy file	Tensor	Magic string mismatch, unsupported version, fortran order, I/O error
read_npz	File path to .npz file	Vec<(String, Tensor)>	Zip errors, header parse errors, fortran order
write_npy	&self + output file path	()	Unsupported kind, I/O error
write_npz	Named tensor pairs + path	()	Unsupported kind, I/O error

NumPy Descriptor	tch Kind
<f2, e	Half
<f4, f	Float
<f8, d	Double
<i4, i	Int
<i8, q	Int64
<i2, h	Int16
<i1, b	Int8
<u1, B	Uint8
?, b1	Bool
<c8, F	ComplexFloat
<c16, D	ComplexDouble

Usage Examples

use tch::{Tensor, Kind, Device};

// Write a tensor to .npy format
let t = Tensor::randn([100, 64], (Kind::Float, Device::Cpu));
t.write_npy("embeddings.npy").unwrap();

// Read it back
let loaded = Tensor::read_npy("embeddings.npy").unwrap();
assert_eq!(loaded.size(), vec![100, 64]);

// Write multiple named tensors to .npz
let weights = Tensor::randn([32, 16], (Kind::Float, Device::Cpu));
let biases = Tensor::zeros([32], (Kind::Float, Device::Cpu));
Tensor::write_npz(
    &[("weights", &weights), ("biases", &biases)],
    "model.npz",
).unwrap();

// Read named tensors from .npz
let named_tensors = Tensor::read_npz("model.npz").unwrap();
for (name, tensor) in &named_tensors {
    println!("{}: {:?}", name, tensor.size());
}

Related Pages

• Principle:LaurentMazare_Tch_rs_NumPy_Format_IO