tensor_file.bend
stelliferous@0.0.3.0 · 495 lines · 0 definitions
1# Raw little-endian FP32 files and NumPy .npy files through the official File
2# effects, and the reader that other formats share: a byte range of an open
3# file, decoded element by element, written straight into a tensor's storage.
4import Base
5import stelliferous@0.0.3.0/storage_bytes.bend as Bytes
6import stelliferous@0.0.3.0/storage_element_format.bend as Formats
7import stelliferous@0.0.3.0/math_fp32.bend as Fp32
8import stelliferous@0.0.3.0/storage_buffer.bend as Buffers
9import stelliferous@0.0.3.0/tensor.bend as Tensors
10import stelliferous@0.0.3.0/tensor_layout.bend as Layout
11
12def Tensor() -> Type:
13 Tensors.Tensor<F32>
14
15def Read() -> Type:
16 File & Result<&1,&1,U32 & String,List<&2,U32>>
17
18def Written() -> Type:
19 File & Result<&1,&1,U32 & String,Unit>
20
21def Sized() -> Type:
22 File & Result<&1,&1,U32 & String,U32>
23
24# FP32 values per File effect; a read takes four times as many bytes.
25def chunk() -> U32:
26 16384
27
28def chunk_bytes() -> U32:
29 (chunk() * 4 : U32)
30
31# Reading: bytes gather into elements, least significant first. The values of
32# one read are written into the buffer before the next read, so no list
33# outgrows a read.
34
35# An element in progress: its first four bytes, the bytes above them, the
36# weight of the next byte within its word, and how many bytes arrived.
37type Partial is Data:
38 Partial{low: U32,high: U32,weight: U32,filled: U32}
39
40type Words is Data:
41 Words{values: List<&2,F32>,count: U32,partial: Partial}
42
43# The buffer of the tensor, the values found so far, and an element in progress.
44type Target is Type:
45 Target{buffer: Buffers.Buffer<F32>,found: U32,partial: Partial}
46
47type Reading is Type:
48 More{file: File,position: U32,remaining: U32,target: Target}
49 End{file: File,target: Target,complete: Bool}
50
51type Outcome is Type:
52 Outcome{file: File,target: Target,complete: Bool}
53
54def fresh_partial() -> Partial:
55 Partial{0,0,1,0}
56
57def gathered(lower: Bool,done: Bool,wrap: Bool,+byte: U32,format: Formats.Format,values: List<&2,F32>,+count: U32,+low: U32,+high: U32,+weight: U32,+filled: U32) -> Words:
58 match lower done wrap:
59 case True{} True{} wrap: Words{Formats.widen(format,(low + byte * weight : U32),high) <> values,(count + 1 : U32),fresh_partial()}
60 case False{} True{} wrap: Words{Formats.widen(format,low,(high + byte * weight : U32)) <> values,(count + 1 : U32),fresh_partial()}
61 case True{} False{} True{}: Words{values,count,Partial{(low + byte * weight : U32),high,1,filled}}
62 case True{} False{} False{}: Words{values,count,Partial{(low + byte * weight : U32),high,(weight * 256 : U32),filled}}
63 case False{} False{} wrap: Words{values,count,Partial{low,(high + byte * weight : U32),(weight * 256 : U32),filled}}
64
65def push(+byte: U32,+format: Formats.Format,words: Words) -> Words:
66 Words{values,+count,partial} = words
67 Partial{+low,+high,+weight,+filled} = partial
68 gathered(U32.is_lt(filled,4),U32.is_eq((filled + 1 : U32),Formats.width(format)),U32.is_eq(weight,16777216),byte,format,values,count,low,high,weight,(filled + 1 : U32))
69
70def trickle(bytes: List<&2,U32>,+format: Formats.Format,words: Words) -> Words:
71 match bytes:
72 case []: words
73 case byte <> rest: trickle(rest,format,push(byte,format,words))
74
75# Whole elements at a time while the bytes allow; the values are newest first.
76def pairs(bytes: List<&2,U32>,+format: Formats.Format,values: List<&2,F32>,+count: U32) -> Words:
77 match bytes:
78 case Con{first,Con{second,rest}}:
79 pairs(rest,format,Formats.widen(format,(first + second * 256 : U32),0) <> values,(count + 1 : U32))
80 case rest: trickle(rest,format,Words{values,count,fresh_partial()})
81
82def quads(bytes: List<&2,U32>,+format: Formats.Format,values: List<&2,F32>,+count: U32) -> Words:
83 match bytes:
84 case Con{first,Con{second,Con{third,Con{fourth,rest}}}}:
85 quads(rest,format,Formats.widen(format,(first + second * 256 + third * 65536 + fourth * 16777216 : U32),0) <> values,(count + 1 : U32))
86 case rest: trickle(rest,format,Words{values,count,fresh_partial()})
87
88def octets(bytes: List<&2,U32>,+format: Formats.Format,values: List<&2,F32>,+count: U32) -> Words:
89 match bytes:
90 case Con{first,Con{second,Con{third,Con{fourth,Con{fifth,Con{sixth,Con{seventh,Con{eighth,rest}}}}}}}}:
91 octets(rest,format,Formats.widen(format,(first + second * 256 + third * 65536 + fourth * 16777216 : U32),
92 (fifth + sixth * 256 + seventh * 65536 + eighth * 16777216 : U32)) <> values,(count + 1 : U32))
93 case rest: trickle(rest,format,Words{values,count,fresh_partial()})
94
95def by_width(+width: U32,bytes: List<&2,U32>,+format: Formats.Format) -> Words:
96 match width:
97 case 2: pairs(bytes,format,[],0)
98 case 4: quads(bytes,format,[],0)
99 case 8: octets(bytes,format,[],0)
100 case _: trickle(bytes,format,Words{[],0,fresh_partial()})
101
102def words(whole: Bool,bytes: List<&2,U32>,+format: Formats.Format,partial: Partial) -> Words:
103 match whole:
104 case True{}: by_width(Formats.width(format),bytes,format)
105 case False{}: trickle(bytes,format,Words{[],0,partial})
106
107# Values beyond the count of the shape are counted and not written.
108def placed(read: Words,buffer: Buffers.Buffer<F32>,+found: U32,+count: U32) -> Target:
109 Words{values,+added,partial} = read
110 Target{Buffers.write_at(~F32,List.take(&2,F32,List.reverse(&2,F32,values),U32.to_nat((count - U32.min(found,count) : U32))),U32.min(found,count),buffer),
111 (found + added : U32),partial}
112
113def fill(bytes: List<&2,U32>,target: Target,+format: Formats.Format,+count: U32) -> Target:
114 Target{buffer,+found,partial} = target
115 Partial{low,high,weight,+filled} = partial
116 placed(words(U32.is_zero(filled),bytes,format,Partial{low,high,weight,filled}),buffer,found,count)
117
118def ranged(exhausted: Bool,file: File,+position: U32,+remaining: U32,target: Target) -> Reading:
119 match exhausted:
120 case True{}: End{file,target,True{}}
121 case False{}: More{file,position,remaining,target}
122
123# A read that ends early leaves the count short, which the count check reports.
124def absorbed(bytes: List<&2,U32>,file: File,+position: U32,+remaining: U32,target: Target,+format: Formats.Format,+count: U32) -> Reading:
125 match bytes:
126 case []: End{file,target,True{}}
127 case Con{+byte,+rest}:
128 +length = U32.min(U32.from_nat(List.length(&2,U32,byte <> rest)),remaining)
129 ranged(U32.is_eq(length,remaining),file,(position + length : U32),(remaining - length : U32),fill(byte <> rest,target,format,count))
130
131def received(result: Result<&1,&1,U32 & String,List<&2,U32>>,file: File,+position: U32,+remaining: U32,target: Target,+format: Formats.Format,+count: U32) -> Reading:
132 match result:
133 case Fail{error}: End{file,target,False{}}
134 case Done{bytes}: absorbed(bytes,file,position,remaining,target,format,count)
135
136def taken(read: Read(),+position: U32,+remaining: U32,target: Target,+format: Formats.Format,+count: U32) -> Reading:
137 (file,result) = read
138 received(result,file,position,remaining,target,format,count)
139
140def read_all(fuel: Nat,reading: Reading,+format: Formats.Format,+count: U32) -> IO(Outcome):
141 match fuel reading:
142 case 0n More{file,position,remaining,target}: IO.pure(Outcome,Outcome{file,target,False{}})
143 case 0n End{file,target,complete}: IO.pure(Outcome,Outcome{file,target,complete})
144 case 1n+rest End{file,target,complete}: IO.pure(Outcome,Outcome{file,target,complete})
145 case 1n+rest More{file,+position,+remaining,target}:
146 do IO<Outcome>:
147 read : Read() <- File.read_at(file,position,U32.min(remaining,chunk_bytes()))
148 read_all(rest,taken(read,position,remaining,target,format,count),format,count)
149
150def fresh(+count: U32) -> Target:
151 Target{Tensors.buffer(~F32,0.0,count),0,fresh_partial()}
152
153def fuel(+length: U32) -> Nat:
154 U32.to_nat((length / chunk_bytes() + 2 : U32))
155
156# The elements stored from the position on, as many bytes as the count of the
157# format takes; the file stays open.
158def read_range(file: File,+position: U32,+length: U32,+format: Formats.Format,+count: U32) -> IO(Outcome):
159 read_all(fuel(length),ranged(U32.is_zero(length),file,position,length,fresh(count)),format,count)
160
161def fault(operation: String,path: String,detail: String,shape: List<&2,U32>) -> Tensor():
162 Tensors.failed(~F32,operation,Tensors.FileFault{path,detail},shape,0.0,1)
163
164def sized(matches: Bool,operation: String,path: String,+shape: List<&2,U32>,buffer: Buffers.Buffer<F32>) -> Tensor():
165 match matches:
166 case True{}: Tensors.from_storage(~F32,shape,buffer,0.0)
167 case False{}: fault(operation,path,"does not hold the values of the shape",shape)
168
169def whole(aligned: Bool,operation: String,path: String,+shape: List<&2,U32>,buffer: Buffers.Buffer<F32>,+found: U32,+count: U32) -> Tensor():
170 match aligned:
171 case False{}: fault(operation,path,"holds a partial element",shape)
172 case True{}: sized(U32.is_eq(found,count),operation,path,shape,buffer)
173
174def aligned(operation: String,path: String,+shape: List<&2,U32>,target: Target,+count: U32) -> Tensor():
175 Target{buffer,+found,partial} = target
176 Partial{low,high,weight,+filled} = partial
177 whole(U32.is_zero(filled),operation,path,shape,buffer,found,count)
178
179def checked(complete: Bool,operation: String,path: String,+shape: List<&2,U32>,target: Target,+count: U32) -> Tensor():
180 match complete:
181 case False{}: fault(operation,path,"cannot be read",shape)
182 case True{}: aligned(operation,path,shape,target,count)
183
184def decoded(operation: String,path: String,+shape: List<&2,U32>,outcome: Outcome,+count: U32) -> File & Tensor():
185 Outcome{file,target,complete} = outcome
186 (file,checked(complete,operation,path,shape,target,count))
187
188def ranged_count(fits: Bool,file: File,operation: String,path: String,+shape: List<&2,U32>,+position: U32,+length: U32,+format: Formats.Format,+count: U32) -> IO(File & Tensor()):
189 match fits:
190 case False{}: IO.pure(File & Tensor(),(file,fault(operation,path,"holds a tensor beyond the machine capacity",shape)))
191 case True{}:
192 do IO<File & Tensor()>:
193 outcome : Outcome <- read_range(file,position,length,format,count)
194 IO.pure(File & Tensor(),decoded(operation,path,shape,outcome,count))
195
196def ranged_shape(count: Maybe<&2,U32>,file: File,operation: String,path: String,+shape: List<&2,U32>,+position: U32,+format: Formats.Format) -> IO(File & Tensor()):
197 match count:
198 case None{}: IO.pure(File & Tensor(),(file,Tensors.failed(~F32,operation,Tensors.InvalidShape{shape},shape,0.0,1)))
199 case Some{+count}:
200 +length = (count * Formats.width(format) : U32)
201 ranged_count(U32.is_eq((length / Formats.width(format) : U32),count),file,operation,path,shape,position,length,format,count)
202
203# The tensor of the shape whose elements of the format start at the position
204# of an open file, which stays open.
205def load_range(file: File,operation: String,+path: String,+shape: List<&2,U32>,+position: U32,+format: Formats.Format) -> IO(File & Tensor()):
206 ranged_shape(Tensors.count_of(shape),file,operation,path,shape,position,format)
207
208def closed(pair: File & Tensor()) -> IO(Tensor()):
209 (file,tensor) = pair
210 do IO<Tensor()>:
211 File.close(file)
212 IO.pure(Tensor(),tensor)
213
214def refused(file: File,tensor: Tensor()) -> IO(Tensor()):
215 do IO<Tensor()>:
216 File.close(file)
217 IO.pure(Tensor(),tensor)
218
219# A raw file holds exactly the bytes of its shape: a size that is not a whole
220# number of values, or another number of them, is refused before reading.
221def measured(aligned: Bool,matches: Bool,file: File,operation: String,+path: String,+shape: List<&2,U32>,+position: U32,+count: U32) -> IO(Tensor()):
222 match aligned matches:
223 case False{} matches: refused(file,fault(operation,path,"is not a whole number of FP32 values",shape))
224 case True{} False{}: refused(file,fault(operation,path,"does not hold the values of the shape",shape))
225 case True{} True{}:
226 do IO<Tensor()>:
227 pair : File & Tensor() <- load_range(file,operation,path,shape,position,Formats.Fp32{})
228 closed(pair)
229
230def with_size(result: Result<&1,&1,U32 & String,U32>,file: File,operation: String,+path: String,+shape: List<&2,U32>,+position: U32,+count: U32) -> IO(Tensor()):
231 match result:
232 case Fail{error}: refused(file,fault(operation,path,"cannot be measured",shape))
233 case Done{+size}:
234 measured(U32.is_zero(((size - U32.min(size,position)) % 4 : U32)),U32.is_le(position,size) && U32.is_eq(((size - U32.min(size,position)) / 4 : U32),count),file,operation,path,shape,position,count)
235
236def sized_file(sized: Sized(),operation: String,+path: String,+shape: List<&2,U32>,+position: U32,+count: U32) -> IO(Tensor()):
237 (file,result) = sized
238 with_size(result,file,operation,path,shape,position,count)
239
240def fp32_values(file: File,operation: String,+path: String,+shape: List<&2,U32>,+position: U32,+count: U32) -> IO(Tensor()):
241 do IO<Tensor()>:
242 sized : Sized() <- File.size(file)
243 sized_file(sized,operation,path,shape,position,count)
244
245def opened(result: Result<&1,&1,U32 & String,File>,+path: String,+shape: List<&2,U32>,+count: U32) -> IO(Tensor()):
246 match result:
247 case Fail{error}: IO.pure(Tensor(),fault("load",path,"cannot be opened",shape))
248 case Done{file}: fp32_values(file,"load",path,shape,0,count)
249
250def counted(count: Maybe<&2,U32>,+path: String,+shape: List<&2,U32>) -> IO(Tensor()):
251 match count:
252 case None{}: IO.pure(Tensor(),Tensors.failed(~F32,"load",Tensors.InvalidShape{shape},shape,0.0,1))
253 case Some{+count}:
254 do IO<Tensor()>:
255 result : Result<&1,&1,U32 & String,File> <- File.open(path,"r")
256 opened(result,path,shape,count)
257
258# The file holds the values of the shape in row-major order and nothing else.
259def load(+path: String,+shape: List<&2,U32>) -> IO(Tensor()):
260 counted(Tensors.count_of(shape),path,shape)
261
262# The NumPy .npy header: magic, version, header length, then a text such as
263# {'descr': '<f4', 'fortran_order': False, 'shape': (16, 3, 6, 6), }
264
265# The shape and the position of the first value.
266type Header is Data:
267 Header{shape: List<&2,U32>,data: U32}
268
269def magic() -> List<&2,U32>:
270 [147,78,85,77,80,89]
271
272def contains(text: List<&2,U32>,+pattern: List<&2,U32>) -> Bool:
273 match text:
274 case Nil{}: False{}
275 case Con{+byte,+rest}: Bytes.starts(pattern,byte <> rest) || contains(rest,pattern)
276
277# The text after the first occurrence of the pattern.
278def after(text: List<&2,U32>,+pattern: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
279 match text:
280 case Nil{}: None{}
281 case Con{+byte,+rest}:
282 Bool.pick(Maybe<&2,List<&2,U32>>,Bytes.starts(pattern,byte <> rest),Some{List.drop(&2,U32,byte <> rest,List.length(&2,U32,pattern))},after(rest,pattern))
283
284def pushed(digits: Bool,current: U32,shape: List<&2,U32>) -> List<&2,U32>:
285 match digits:
286 case True{}: current <> shape
287 case False{}: shape
288
289def closed_shape(valid: Bool,shape: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
290 match valid:
291 case True{}: Some{List.reverse(&2,U32,shape)}
292 case False{}: None{}
293
294# Extents up to the closing parenthesis: digits, commas and spaces.
295def extents(text: List<&2,U32>,+current: U32,digits: Bool,valid: Bool,shape: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
296 match text:
297 case Nil{}: None{}
298 case Con{+byte,rest}:
299 match byte:
300 case 41: closed_shape(valid,pushed(digits,current,shape))
301 case 44: extents(rest,0,False{},valid,pushed(digits,current,shape))
302 case 32: extents(rest,current,digits,valid,shape)
303 case _: extents(rest,(current * 10 + byte - 48 : U32),True{},valid && U32.is_le(48,byte) && U32.is_le(byte,57),shape)
304
305def shaped(shape: Maybe<&2,List<&2,U32>>,+data: U32) -> Maybe<&2,Header>:
306 match shape:
307 case None{}: None{}
308 case Some{shape}: Some{Header{shape,data}}
309
310def typed(supported: Bool,text: Maybe<&2,List<&2,U32>>,+data: U32) -> Maybe<&2,Header>:
311 match supported text:
312 case False{} text: None{}
313 case True{} None{}: None{}
314 case True{} Some{text}: shaped(extents(text,0,False{},True{},[]),data)
315
316def described(+text: List<&2,U32>,+length: Nat,+data: U32) -> Maybe<&2,Header>:
317 typed(Nat.is_eq(List.length(&2,U32,text),length) && contains(text,Bytes.codes("'descr': '<f4'")) && contains(text,Bytes.codes("'fortran_order': False")),
318 after(text,Bytes.codes("'shape': (")),data)
319
320def region(+bytes: List<&2,U32>,+prefix: U32,+length: U32) -> Maybe<&2,Header>:
321 described(List.take(&2,U32,List.drop(&2,U32,bytes,U32.to_nat(prefix)),U32.to_nat(length)),U32.to_nat(length),(prefix + length : U32))
322
323# Version 1 has a two-byte header length; later versions have four bytes.
324def versioned(major: U32,+bytes: List<&2,U32>) -> Maybe<&2,Header>:
325 match major:
326 case 1: region(bytes,10,Bytes.u16_at(bytes,8n))
327 case _: region(bytes,12,Bytes.u32_at(bytes,8n))
328
329def marked(valid: Bool,+bytes: List<&2,U32>) -> Maybe<&2,Header>:
330 match valid:
331 case False{}: None{}
332 case True{}: versioned(Bytes.at(bytes,6n),bytes)
333
334def header(+bytes: List<&2,U32>) -> Maybe<&2,Header>:
335 marked(Bytes.starts(magic(),bytes),bytes)
336
337def npy_counted(count: Maybe<&2,U32>,+shape: List<&2,U32>,+data: U32,file: File,+path: String) -> IO(Tensor()):
338 match count:
339 case None{}: refused(file,Tensors.failed(~F32,"load_npy",Tensors.InvalidShape{shape},shape,0.0,1))
340 case Some{+count}: fp32_values(file,"load_npy",path,shape,data,count)
341
342def npy_parsed(parsed: Maybe<&2,Header>,file: File,+path: String) -> IO(Tensor()):
343 match parsed:
344 case None{}: refused(file,fault("load_npy",path,"is not a little-endian FP32 .npy array in C order",[]))
345 case Some{Header{+shape,data}}: npy_counted(Tensors.count_of(shape),shape,data,file,path)
346
347def npy_received(result: Result<&1,&1,U32 & String,List<&2,U32>>,file: File,+path: String) -> IO(Tensor()):
348 match result:
349 case Fail{error}: refused(file,fault("load_npy",path,"cannot be read",[]))
350 case Done{bytes}: npy_parsed(header(bytes),file,path)
351
352def npy_headed(read: Read(),+path: String) -> IO(Tensor()):
353 (file,result) = read
354 npy_received(result,file,path)
355
356def npy_opened(result: Result<&1,&1,U32 & String,File>,+path: String) -> IO(Tensor()):
357 match result:
358 case Fail{error}: IO.pure(Tensor(),fault("load_npy",path,"cannot be opened",[]))
359 case Done{file}:
360 do IO<Tensor()>:
361 read : Read() <- File.read_at(file,0,chunk_bytes())
362 npy_headed(read,path)
363
364# An FP32 array saved by NumPy, with its shape.
365def load_npy(+path: String) -> IO(Tensor()):
366 do IO<Tensor()>:
367 result : Result<&1,&1,U32 & String,File> <- File.open(path,"r")
368 npy_opened(result,path)
369
370# Writing
371
372type Writing is Type:
373 Pending{file: File,storage: Array<F32>,offset: U32,remaining: Nat}
374 Finished{file: File,complete: Bool}
375
376def bytes(+word: U32,rest: List<&2,U32>) -> List<&2,U32>:
377 (word % 256 : U32) <> (word / 256 % 256 : U32) <> (word / 65536 % 256 : U32) <> (word / 16777216 : U32) <> rest
378
379def encode(count: Nat,values: List<&2,F32>) -> List<&2,U32>:
380 match count values:
381 case 0n values: []
382 case 1n+rest Nil{}: []
383 case 1n+rest Con{value,more}: bytes(Fp32.to_bits(value),encode(rest,more))
384
385# The cursor retains the array; only the current chunk becomes a list.
386def read_chunk(count: Nat,+offset: U32,storage: Array<F32>,values: List<&2,F32>) -> Array<F32> & List<&2,F32>:
387 match count:
388 case 0n: (storage,List.reverse(&2,F32,values))
389 case 1n+rest:
390 Buffers.bind_pair(Array<F32>,F32,Array<F32> & List<&2,F32>,Array.get(F32,storage,offset),storage => value =>
391 read_chunk(rest,(offset + 1 : U32),storage,value <> values))
392
393def pending(remaining: Nat,file: File,storage: Array<F32>,offset: U32) -> Writing:
394 match remaining:
395 case 0n: Finished{file,True{}}
396 case 1n+rest: Pending{file,storage,offset,1n+rest}
397
398def continued(result: Result<&1,&1,U32 & String,Unit>,file: File,storage: Array<F32>,remaining: Nat,offset: U32) -> Writing:
399 match result:
400 case Fail{error}: Finished{file,False{}}
401 case Done{unit}: pending(remaining,file,storage,offset)
402
403def after_write(written: Written(),storage: Array<F32>,remaining: Nat,offset: U32) -> Writing:
404 (file,result) = written
405 continued(result,file,storage,remaining,offset)
406
407def write_chunk(file: File,read: Array<F32> & List<&2,F32>,+count: Nat,remaining: Nat,offset: U32) -> IO(Writing):
408 (storage,values) = read
409 do IO<Writing>:
410 written : Written() <- File.write_bytes(file,encode(count,values))
411 IO.pure(Writing,after_write(written,storage,Nat.sub(remaining,count),(offset + U32.from_nat(count) : U32)))
412
413def finished(file: File,complete: Bool) -> IO(Bool):
414 do IO<Bool>:
415 File.close(file)
416 IO.pure(Bool,complete)
417
418def write_all(fuel: Nat,writing: Writing) -> IO(Bool):
419 match fuel writing:
420 case 0n Pending{file,storage,offset,remaining}: finished(file,False{})
421 case 0n Finished{file,complete}: finished(file,complete)
422 case 1n+rest Finished{file,complete}: finished(file,complete)
423 case 1n+rest Pending{file,storage,+offset,+remaining}:
424 +count = Nat.min(U32.to_nat(chunk()),remaining)
425 do IO<Bool>:
426 next : Writing <- write_chunk(file,read_chunk(count,offset,storage,[]),count,remaining,offset)
427 write_all(rest,next)
428
429def unsaved(pair: Tensor() & List<&2,U32>,operation: String,path: String,detail: String) -> Tensor():
430 (tensor,shape) = pair
431 fault(operation,path,detail,shape)
432
433def written(complete: Bool,operation: String,path: String,tensor: Tensor()) -> Tensor():
434 match complete:
435 case True{}: tensor
436 case False{}: unsaved(Tensors.shape(~F32,tensor),operation,path,"cannot be written")
437
438# The prefix is written first; a raw file has none.
439def created(result: Result<&1,&1,U32 & String,File>,operation: String,+path: String,tensor: Tensor(),prefix: List<&2,U32>,storage: Array<F32>,+length: Nat) -> IO(Tensor()):
440 match result:
441 case Fail{error}: IO.pure(Tensor(),unsaved(Tensors.shape(~F32,tensor),operation,path,"cannot be created"))
442 case Done{file}:
443 do IO<Tensor()>:
444 head : Written() <- File.write_bytes(file,prefix)
445 complete : Bool <- write_all(Nat.add(length,2n),after_write(head,storage,length,0))
446 IO.pure(Tensor(),written(complete,operation,path,tensor))
447
448def saving(operation: String,+path: String,tensor: Tensor(),prefix: List<&2,U32>,buffer: Result<&2,&1,Tensors.Problem,Buffers.Buffer<F32>>) -> IO(Tensor()):
449 match buffer:
450 case Fail{problem}: IO.pure(Tensor(),tensor)
451 case Done{Buffers.Buffer{storage,length,certificate}}:
452 do IO<Tensor()>:
453 result : Result<&1,&1,U32 & String,File> <- File.open(path,"w")
454 created(result,operation,path,tensor,prefix,storage,length)
455
456# Writes the values in row-major order and returns the tensor. A failed tensor
457# writes nothing and keeps its problem.
458def save(+path: String,tensor: Tensor()) -> IO(Tensor()):
459 Tensors.then(~Tensor(),~Tensor(),~IO(Tensor()),Layout.copy(~F32,tensor),kept => copy => saving("save",path,kept,[],Tensors.into_storage(~F32,copy)))
460
461def joined(shape: List<&2,U32>) -> String:
462 match shape:
463 case []: ""
464 case [extent]: U32.show(extent)
465 case extent <> rest: U32.show(extent) ++ ", " ++ joined(rest)
466
467# NumPy writes (), (3,) and (2, 3).
468def tuple(shape: List<&2,U32>) -> String:
469 match shape:
470 case []: "()"
471 case [extent]: "(" ++ U32.show(extent) ++ ",)"
472 case extent <> rest: "(" ++ joined(extent <> rest) ++ ")"
473
474def spaces(count: Nat) -> List<&2,U32>:
475 match count:
476 case 0n: [10]
477 case 1n+rest: 32 <> spaces(rest)
478
479# Version 1: the header ends in a newline and the data starts at a multiple of 64 bytes.
480def framed(+text: List<&2,U32>,+length: U32) -> List<&2,U32>:
481 +padding = ((64 - (11 + length) % 64) % 64 : U32)
482 +total = (length + padding + 1 : U32)
483 List.append(&2,U32,magic(),1 <> 0 <> (total % 256 : U32) <> (total / 256 : U32) <> List.append(&2,U32,text,spaces(U32.to_nat(padding))))
484
485def npy_header(shape: List<&2,U32>) -> List<&2,U32>:
486 +text = Bytes.codes("{'descr': '<f4', 'fortran_order': False, 'shape': " ++ tuple(shape) ++ ", }")
487 framed(text,U32.from_nat(List.length(&2,U32,text)))
488
489def npy_saving(+path: String,pair: Tensor() & Tensor(),+shape: List<&2,U32>) -> IO(Tensor()):
490 (kept,copy) = pair
491 saving("save_npy",path,kept,npy_header(shape),Tensors.into_storage(~F32,copy))
492
493# Writes a .npy file that NumPy loads with the tensor's shape.
494def save_npy(+path: String,tensor: Tensor()) -> IO(Tensor()):
495 Tensors.then(~Tensor(),~List<&2,U32>,~IO(Tensor()),Tensors.shape(~F32,tensor),tensor => shape => npy_saving(path,Layout.copy(~F32,tensor),shape))