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/// Handle that reads a tensor's bytes.
44type Read = Arc<dyn Fn() -> io::Result<Vec<u8>> + Send + Sync>;
45
46/// A tensor in a checkpoint, with its bytes not yet read.
47///
48/// Opening a checkpoint parses its metadata only. [`read`](Self::read) is what produces the
49/// tensor's bytes, so a checkpoint can be inspected without loading it, or loaded one tensor
50/// at a time. ZIP and legacy containers read the file at that point; a TAR container's
51/// storages are held in memory from open.
52///
53/// The element type and shape are fixed at parse time, when the reader checks that the
54/// tensor fits its storage, so they are read-only: [`byte_len`](Self::byte_len) is what
55/// [`read`](Self::read) returns because nothing can move them apart.
56#[derive(Clone)]
57pub struct Tensor {
58 /// Dotted path of the tensor through the checkpoint's dictionaries
59 /// (`"encoder.0.weight"`). Free to rename; only error messages look at it.
60 pub name: String,
61 dtype: DType,
62 shape: Vec<usize>,
63 read: Read,
64 /// Reads the transpose of a 2-D tensor the checkpoint stores transposed.
65 transposed: Option<Read>,
66}
67
68impl Tensor {
69 pub(crate) fn new(
70 name: String,
71 dtype: DType,
72 shape: Vec<usize>,
73 read: impl Fn() -> io::Result<Vec<u8>> + Send + Sync + 'static,
74 ) -> Self {
75 Self {
76 name,
77 dtype,
78 shape,
79 read: Arc::new(read),
80 transposed: None,
81 }
82 }
83
84 /// Attach a handle that reads the tensor's transpose as a straight copy.
85 pub(crate) fn with_transposed(
86 mut self,
87 read: impl Fn() -> io::Result<Vec<u8>> + Send + Sync + 'static,
88 ) -> Self {
89 self.transposed = Some(Arc::new(read));
90 self
91 }
92
93 /// Element type.
94 pub fn dtype(&self) -> DType {
95 self.dtype
96 }
97
98 /// Extents, outermost first. Empty for a scalar.
99 pub fn shape(&self) -> &[usize] {
100 &self.shape
101 }
102
103 /// Number of elements.
104 pub fn num_elements(&self) -> usize {
105 self.shape.iter().product()
106 }
107
108 /// Number of bytes [`read`](Self::read) returns.
109 pub fn byte_len(&self) -> usize {
110 self.num_elements() * self.dtype.size()
111 }
112
113 /// Read the tensor's bytes from the checkpoint.
114 ///
115 /// The bytes are contiguous in row-major order and native-endian, whatever layout the
116 /// file stored the tensor in: a strided or offset view is gathered, and a `bool` tensor
117 /// holds only 0 and 1. Nothing is cached, so each call reads the storage again.
118 ///
119 /// Reads from different threads run at once: a ZIP read holds the archive lock only to
120 /// locate its entry. A compressed entry, which `torch.save` does not write, streams
121 /// under the lock instead, as does every entry on a target other than Unix.
122 ///
123 /// # Errors
124 ///
125 /// A storage that is missing, corrupt or unreadable by its container fails with
126 /// [`io::ErrorKind::InvalidData`] (a ZIP storage's checksum is verified when the tensor
127 /// spans all of it, unless the file was saved without checksums, as
128 /// `compute_crc32=False` does), and one shorter than the pickle declared with
129 /// [`io::ErrorKind::UnexpectedEof`]: the file disagrees with itself. Any other kind is
130 /// the operating system's, and means the file could not be read at all. Every error
131 /// names the tensor and its storage.
132 pub fn read(&self) -> io::Result<Vec<u8>> {
133 (self.read)().map_err(|err| named(&self.name, err))
134 }
135
136 /// The transpose of a 2-D tensor the checkpoint stores transposed.
137 ///
138 /// A `.t()` view, which is how `nn.Linear` weights are often saved, is stored as the
139 /// row-major bytes of its transpose. [`read`](Self::read) has to gather those into the
140 /// view's own order, while the tensor returned here, with the reversed shape, reads them
141 /// as a straight copy. A consumer that wants the transpose anyway can skip both passes.
142 ///
143 /// `None` for every other tensor, including one whose transpose would itself need
144 /// gathering. The returned tensor has the same name and reads and fails as
145 /// [`read`](Self::read) does; its own `transposed` is `None`.
146 #[must_use]
147 pub fn transposed(&self) -> Option<Tensor> {
148 let read = self.transposed.clone()?;
149 Some(Self {
150 name: self.name.clone(),
151 dtype: self.dtype,
152 shape: self.shape.iter().rev().copied().collect(),
153 read,
154 transposed: None,
155 })
156 }
157
158 /// Split the tensor into its metadata and its read handle.
159 ///
160 /// For a caller that keeps the metadata elsewhere and wants only the handle, without a
161 /// second copy of the name and shape held alive with it. The handle reads and fails
162 /// exactly as [`read`](Self::read) does.
163 pub fn into_parts(
164 self,
165 ) -> (
166 String,
167 DType,
168 Vec<usize>,
169 impl Fn() -> io::Result<Vec<u8>> + Send + Sync + 'static,
170 ) {
171 let Self {
172 name,
173 dtype,
174 shape,
175 read,
176 ..
177 } = self;
178 let context = name.clone();
179
180 (name, dtype, shape, move || {
181 read().map_err(|err| named(&context, err))
182 })
183 }
184}
185
186impl fmt::Debug for Tensor {
187 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188 f.debug_struct("Tensor")
189 .field("name", &self.name)
190 .field("dtype", &self.dtype)
191 .field("shape", &self.shape)
192 .finish_non_exhaustive()
193 }
194}
195
196/// Put the tensor's name in front of a read error, keeping its kind.
197///
198/// The reader builds a tensor before it knows the name (that is assembled from the dict
199/// path afterwards), so the name is added here rather than inside the handle.
200fn named(name: &str, err: io::Error) -> io::Error {
201 if name.is_empty() {
202 return err;
203 }
204 io::Error::new(err.kind(), format!("tensor '{name}': {err}"))
205}