pytorch_reader/
tensor.rs

1//! What the reader produces for each tensor in a checkpoint.
2
3use std::fmt;
4use std::io;
5use std::sync::Arc;
6
7/// Element type of a tensor.
8///
9/// The types `torch.save` writes as a plain storage. Complex, quantized and sparse tensors
10/// are refused at parse time rather than given a variant here. Deliberately exhaustive: a
11/// converter matching on it should fail to compile, not fall through, when a variant is
12/// added, so adding one is a breaking change.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum DType {
15    F64,
16    F32,
17    F16,
18    BF16,
19    I64,
20    I32,
21    I16,
22    I8,
23    U64,
24    U32,
25    U16,
26    U8,
27    /// One byte per element, holding 0 or 1.
28    Bool,
29}
30
31impl DType {
32    /// Size of one element in bytes.
33    pub const fn size(self) -> usize {
34        match self {
35            DType::F64 | DType::I64 | DType::U64 => 8,
36            DType::F32 | DType::I32 | DType::U32 => 4,
37            DType::F16 | DType::BF16 | DType::I16 | DType::U16 => 2,
38            DType::I8 | DType::U8 | DType::Bool => 1,
39        }
40    }
41}
42
43/// A tensor in a checkpoint, with its bytes not yet read.
44///
45/// Opening a checkpoint parses its metadata only. [`read`](Self::read) is what produces the
46/// tensor's bytes, so a checkpoint can be inspected without loading it, or loaded one tensor
47/// at a time. ZIP and legacy containers read the file at that point; a TAR container's
48/// storages are held in memory from open.
49///
50/// The element type and shape are fixed at parse time, when the reader checks that the
51/// tensor fits its storage, so they are read-only: [`byte_len`](Self::byte_len) is what
52/// [`read`](Self::read) returns because nothing can move them apart.
53#[derive(Clone)]
54pub struct Tensor {
55    /// Dotted path of the tensor through the checkpoint's dictionaries
56    /// (`"encoder.0.weight"`). Free to rename; only error messages look at it.
57    pub name: String,
58    dtype: DType,
59    shape: Vec<usize>,
60    read: Arc<dyn Fn() -> io::Result<Vec<u8>> + Send + Sync>,
61}
62
63impl Tensor {
64    pub(crate) fn new(
65        name: String,
66        dtype: DType,
67        shape: Vec<usize>,
68        read: impl Fn() -> io::Result<Vec<u8>> + Send + Sync + 'static,
69    ) -> Self {
70        Self {
71            name,
72            dtype,
73            shape,
74            read: Arc::new(read),
75        }
76    }
77
78    /// Element type.
79    pub fn dtype(&self) -> DType {
80        self.dtype
81    }
82
83    /// Extents, outermost first. Empty for a scalar.
84    pub fn shape(&self) -> &[usize] {
85        &self.shape
86    }
87
88    /// Number of elements.
89    pub fn num_elements(&self) -> usize {
90        self.shape.iter().product()
91    }
92
93    /// Number of bytes [`read`](Self::read) returns.
94    pub fn byte_len(&self) -> usize {
95        self.num_elements() * self.dtype.size()
96    }
97
98    /// Read the tensor's bytes from the checkpoint.
99    ///
100    /// The bytes are contiguous in row-major order and native-endian, whatever layout the
101    /// file stored the tensor in: a strided or offset view is gathered, and a `bool` tensor
102    /// holds only 0 and 1. Nothing is cached, so each call reads the storage again.
103    ///
104    /// Reads from different threads run at once: a ZIP read holds the archive lock only to
105    /// locate its entry. A compressed entry, which `torch.save` does not write, streams
106    /// under the lock instead, as does every entry on a target other than Unix.
107    ///
108    /// # Errors
109    ///
110    /// A storage that is missing, corrupt or unreadable by its container fails with
111    /// [`io::ErrorKind::InvalidData`] (a ZIP storage's checksum is verified when the tensor
112    /// spans all of it, unless the file was saved without checksums, as
113    /// `compute_crc32=False` does), and one shorter than the pickle declared with
114    /// [`io::ErrorKind::UnexpectedEof`]: the file disagrees with itself. Any other kind is
115    /// the operating system's, and means the file could not be read at all. Every error
116    /// names the tensor and its storage.
117    pub fn read(&self) -> io::Result<Vec<u8>> {
118        (self.read)().map_err(|err| named(&self.name, err))
119    }
120
121    /// Split the tensor into its metadata and its read handle.
122    ///
123    /// For a caller that keeps the metadata elsewhere and wants only the handle, without a
124    /// second copy of the name and shape held alive with it. The handle reads and fails
125    /// exactly as [`read`](Self::read) does.
126    pub fn into_parts(
127        self,
128    ) -> (
129        String,
130        DType,
131        Vec<usize>,
132        impl Fn() -> io::Result<Vec<u8>> + Send + Sync + 'static,
133    ) {
134        let Self {
135            name,
136            dtype,
137            shape,
138            read,
139        } = self;
140        let context = name.clone();
141
142        (name, dtype, shape, move || {
143            read().map_err(|err| named(&context, err))
144        })
145    }
146}
147
148impl fmt::Debug for Tensor {
149    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
150        f.debug_struct("Tensor")
151            .field("name", &self.name)
152            .field("dtype", &self.dtype)
153            .field("shape", &self.shape)
154            .finish_non_exhaustive()
155    }
156}
157
158/// Put the tensor's name in front of a read error, keeping its kind.
159///
160/// The reader builds a tensor before it knows the name (that is assembled from the dict
161/// path afterwards), so the name is added here rather than inside the handle.
162fn named(name: &str, err: io::Error) -> io::Error {
163    if name.is_empty() {
164        return err;
165    }
166    io::Error::new(err.kind(), format!("tensor '{name}': {err}"))
167}
pytorch_reader/tensor.rs

pytorch_reader/
tensor.rs
