main.bend

bendtorch-load@1.2.0.0 · 612 lines · 96 definitions

1# Temporary loading-only exception: bend-ml-binary-read is foreign C/JavaScript. No file with laws imports this module.
2# Replace it with native Bend when the language supports dense reads (Renan, 2026-10-17).
3# The reader is bend-ml-binary-read@0.2.0.0, with 64-bit offsets, BF16 files and reads into a cell offset
4import Base
5import bend-ml-binary-read@0.2.0.0/main.bend as R
6import bendtorch@1.2.0.0/main.bend as T
7
8def require(ok: Bool, +message: String) -> IO(Unit):
9  match ok:
10    case True{}: IO.pure(Unit, Unit{})
11    case False{}: IO.die(Unit, 1, message)
12
13# =====================================================================================================
14# W64: a 64-bit byte count or offset as two U32s, hi * 2^32 + lo. Bend (2.0.39) has no 64-bit integer, and a file past
15# 4 GiB needs one (SmolLM2-1.7B's F32 weights, a 360M AdamW checkpoint). Plain Bend: none of these reaches the reader.
16# =====================================================================================================
17
18type W64 is Data:
19  W64{hi: U32, lo: U32}
20
21def W64.of(x: U32) -> W64:
22  W64{0, x}
23
24def W64.of_pair(p: U32 & U32) -> W64:
25  match p:
26    case Tuple{hi, lo}: W64{hi, lo}
27
28def W64.bit(b: Bool) -> U32:
29  match b:
30    case True{}: 1
31    case False{}: 0
32
33def W64.eq(a: W64, b: W64) -> Bool:
34  match a b:
35    case W64{ah, al} W64{bh, bl}: Bool.and(U32.is_eq(ah, bh), U32.is_eq(al, bl))
36
37def W64.add_go(+ah: U32, +bh: U32, +s: U32, +al: U32) -> W64:
38  W64{((ah + bh : U32) + W64.bit(U32.is_lt(s, al)) : U32), s}
39
40# a + b, wrapping at 2^64: the low words' sum carries one when it wrapped (it is then below either term)
41def W64.add(a: W64, b: W64) -> W64:
42  match a b:
43    case W64{ah, +al} W64{bh, bl}: W64.add_go(ah, bh, (al + bl : U32), al)
44
45def W64.mul_go(+h: U32, +k: U32, +p0: U32, +p1: U32) -> W64:
46  W64.add(W64{((h * k : U32) + U32.shrn(p1, 16n) : U32), p0}, W64{0, U32.shln(p1, 16n)})
47
48# a * k for k < 2^16 (enough for bytes a number and decimal digits), wrapping at 2^64: lo = l1 2^16 + l0, so
49# lo * k = (l1 k) 2^16 + l0 k, where both products fit a U32
50def W64.mul_small(a: W64, +k: U32) -> W64:
51  match a:
52    case W64{h, +l}: W64.mul_go(h, k, (U32.and(l, 65535) * k : U32), (U32.shrn(l, 16n) * k : U32))
53
54def W64.div10_go(+qh: U32, +r1: U32, +q1: U32, +l: U32) -> W64 & U32:
55  +t = (U32.shln(r1, 16n) + U32.and(l, 65535) : U32)
56  Tuple{W64{qh, U32.or(U32.shln(q1, 16n), U32.div(t, 10))}, U32.mod(t, 10)}
57
58# (a / 10, a % 10), by long division in 16-bit steps so no partial remainder overflows a U32
59def W64.div10(a: W64) -> W64 & U32:
60  match a:
61    case W64{+h, +l}:
62      +t = (U32.shln(U32.mod(h, 10), 16n) + U32.shrn(l, 16n) : U32)
63      W64.div10_go(U32.div(h, 10), U32.mod(t, 10), U32.div(t, 10), l)
64
65def W64.is_zero(a: W64) -> Bool:
66  match a:
67    case W64{h, l}: Bool.and(U32.is_zero(h), U32.is_zero(l))
68
69def W64.show_go(f: Nat, z: Bool, qr: W64 & U32, acc: String) -> String:
70  match f z qr:
71    case _ True{} Tuple{q, r}: acc
72    case 0n False{} Tuple{q, r}: acc
73    case 1n+g False{} Tuple{+q, r}: W64.show_go(g, W64.is_zero(q), W64.div10(q), U32.show(r) ++ acc)
74
75def W64.show_if(z: Bool, +a: W64) -> String:
76  match z:
77    case True{}: "0"
78    case False{}: W64.show_go(20n, False{}, W64.div10(a), "")
79
80# the decimal digits, as U32.show writes them
81def W64.show(+a: W64) -> String:
82  W64.show_if(W64.is_zero(a), a)
83
84def W64.digits(s: String, f: Nat) -> Bool:
85  match s f:
86    case SNil{} _: True{}
87    case SCon{c, t} 0n: False{}
88    case SCon{c, t} 1n+g: Bool.and(Char.is_digit(c), W64.digits(t, g))
89
90def W64.read_go(s: String, +acc: W64) -> W64:
91  match s:
92    case SNil{}: acc
93    case SCon{Chr{x}, t}: W64.read_go(t, W64.add(W64.mul_small(acc, 10), W64.of(U32.sub(x, 48))))
94
95def W64.read_if(ok: Bool, s: String) -> Maybe<&2, W64>:
96  match ok:
97    case True{}: Some{W64.read_go(s, W64{0, 0})}
98    case False{}: None{}
99
100# A decimal number of 1 to 19 digits (every such number fits 64 bits); anything else is None
101def W64.read(s: String) -> Maybe<&2, W64>:
102  match s:
103    case SNil{}: None{}
104    case SCon{+h, +t}: W64.read_if(W64.digits(SCon{h, t}, 19n), SCon{h, t})
105
106# =====================================================================================================
107# Reading
108# =====================================================================================================
109
110def read_done(+path: String, +n: U32, r: File & (Array<F32> & Result<&1, &1, U32 & String, U32>)) -> IO(Array<F32>):
111  match r:
112    case Tuple{f, Tuple{a, res}}:
113      do IO<Array<F32>>:
114        +got : U32 <- IO.pass(U32, res)
115        File.close(f)
116        require(U32.is_eq(got, n), path ++ ": short read (" ++ U32.show(got) ++ " of " ++ U32.show(n) ++ " numbers)")
117        return a
118
119def read_go(+path: String, +n: U32, f: File, +size: W64) -> IO(Array<F32>):
120  do IO<Array<F32>>:
121    require(W64.eq(size, W64.mul_small(W64.of(n), 4)), path ++ ": expected " ++ W64.show(W64.mul_small(W64.of(n), 4)) ++ " bytes, got " ++ W64.show(size))
122    rr : File & (Array<F32> & Result<&1, &1, U32 & String, U32>) <- R.read_f32_at(f, 0, 0, n, Array.new(F32, T.depth(U32.to_nat(n)), 0.0))
123    read_done(path, n, rr)
124
125def read_sized(+path: String, +n: U32, r: File & Result<&1, &1, U32 & String, U32 & U32>) -> IO(Array<F32>):
126  match r:
127    case Tuple{f, res}:
128      do IO<Array<F32>>:
129        sz : U32 & U32 <- IO.pass(U32 & U32, res)
130        read_go(path, n, f, W64.of_pair(sz))
131
132# A raw, exact-size little-endian F32 file, read once into one source array: up to 2^32 - 1 numbers (16 GiB).
133def read(+path: String, +n: U32) -> IO(Array<F32>):
134  do IO<Array<F32>>:
135    f : File <- IO.try(File, File.open(path, "r"))
136    rr : File & Result<&1, &1, U32 & String, U32 & U32> <- R.size64(f)
137    read_sized(path, n, rr)
138
139# n little-endian F32 numbers at byte `off` of a file of any size (a tensor inside a larger file)
140def read_f32_at(+path: String, off: W64, +n: U32) -> IO(Array<F32>):
141  match off:
142    case W64{hi, lo}:
143      do IO<Array<F32>>:
144        f : File <- IO.try(File, File.open(path, "r"))
145        rr : File & (Array<F32> & Result<&1, &1, U32 & String, U32>) <- R.read_f32_at(f, hi, lo, n, Array.new(F32, T.depth(U32.to_nat(n)), 0.0))
146        read_done(path, n, rr)
147
148# n little-endian BF16 numbers at byte `off` of a file of any size, each widened exactly to F32 (bf16.to(torch.float32))
149def read_bf16_at(+path: String, off: W64, +n: U32) -> IO(Array<F32>):
150  match off:
151    case W64{hi, lo}:
152      do IO<Array<F32>>:
153        f : File <- IO.try(File, File.open(path, "r"))
154        rr : File & (Array<F32> & Result<&1, &1, U32 & String, U32>) <- R.read_bf16_at(f, hi, lo, n, Array.new(F32, T.depth(U32.to_nat(n)), 0.0))
155        read_done(path, n, rr)
156
157def size_done(r: File & Result<&1, &1, U32 & String, U32 & U32>) -> IO(W64):
158  match r:
159    case Tuple{f, res}:
160      do IO<W64>:
161        sz : U32 & U32 <- IO.pass(U32 & U32, res)
162        File.close(f)
163        return W64.of_pair(sz)
164
165# A file's size in bytes, at any size
166def size(+path: String) -> IO(W64):
167  do IO<W64>:
168    f : File <- IO.try(File, File.open(path, "r"))
169    rr : File & Result<&1, &1, U32 & String, U32 & U32> <- R.size64(f)
170    size_done(rr)
171
172def weights_done(-P: Type, +n: U32, r: T.PA<P>) -> IO(P):
173  match r:
174    case T.PA{p, xs, +off}:
175      do IO<P>:
176        require(U32.is_eq(off, n), "Layer.count differs from the loading cursor")
177        return p
178
179# All parameters in Layer.dump order. The source array is released after the model's arrays are filled.
180def weights(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer<P, X, Y, S>, +path: String, p: P) -> IO(P):
181  do IO<P>:
182    +n : U32 = T.Layer.count(P, X, Y, S, m)
183    a : Array<F32> <- read(path, n)
184    weights_done(P, n, T.Layer.load(P, X, Y, S, m, p, a, 0))
185
186# The checkpoint layout is determined by the already initialized optimizer state, without a header:
187# SGD: parameters; momentum: parameters, buffer; Adam/AdamW: parameters, m, v, 1 - b1^t, 1 - b2^t; AdamW with packed
188# BF16 moments: parameters, the moments' tree (main.bend's pt_store), 1 - b1^t, 1 - b2^t, the whole file read.
189def train_read(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer<P, X, Y, S>, +path: String, +base: U32, +copies: U32, +tail: U32, tr: T.Train<P>) -> IO(T.Train<P>):
190  do IO<T.Train<P>>:
191    require(U32.is_le(base, U32.div(U32.sub(4294967295, tail), copies)), "checkpoint exceeds 2^32 - 1 numbers")
192    +n : U32 = ((copies * base : U32) + tail : U32)
193    a : Array<F32> <- read(path, n)
194    return T.Train.restore(~P, ~X, ~Y, ~S, ~m, tr, a)
195
196def cells_done(r: File & Result<&1, &1, U32 & String, U32>) -> IO(U32):
197  match r:
198    case Tuple{f, res}:
199      do IO<U32>:
200        +size : U32 <- IO.pass(U32, res)
201        File.close(f)
202        return U32.div(size, 4)
203
204# the number of F32 cells in a file (a packed-moment checkpoint carries its own sizes: its length is the file's)
205def cells_of(+path: String) -> IO(U32):
206  do IO<U32>:
207    f : File <- IO.try(File, File.open(path, "r"))
208    rr : File & Result<&1, &1, U32 & String, U32> <- File.size(f)
209    cells_done(rr)
210
211def train(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer<P, X, Y, S>, +path: String, tr: T.Train<P>) -> IO(T.Train<P>):
212  match tr:
213    case T.Train{p, g, T.NoMom{}, p1, p2, loss}:
214      train_read(~P, ~X, ~Y, ~S, ~m, path, T.Layer.count(P, X, Y, S, m), 1, 0, T.Train{p, g, T.NoMom{}, p1, p2, loss})
215    case T.Train{p, g, T.OneMom{mm}, p1, p2, loss}:
216      train_read(~P, ~X, ~Y, ~S, ~m, path, T.Layer.count(P, X, Y, S, m), 2, 0, T.Train{p, g, T.OneMom{mm}, p1, p2, loss})
217    case T.Train{p, g, T.TwoMom{mm, vv}, p1, p2, loss}:
218      train_read(~P, ~X, ~Y, ~S, ~m, path, T.Layer.count(P, X, Y, S, m), 3, 2, T.Train{p, g, T.TwoMom{mm, vv}, p1, p2, loss})
219    case T.Train{p, g, T.Packed{t}, p1, p2, loss}:
220      do IO<T.Train<P>>:
221        +n : U32 <- cells_of(path)
222        train_read(~P, ~X, ~Y, ~S, ~m, path, n, 1, 0, T.Train{p, g, T.Packed{t}, p1, p2, loss})
223
224# =====================================================================================================
225# Safetensors (Hugging Face): 8 bytes holding the header's length (little-endian), a JSON header naming each tensor's
226# dtype, shape and byte range (data_offsets, counted from the end of the header), then the data. The header is read as
227# bytes and scanned in plain Bend: a lexer into tokens, then one pass that keeps what a tensor's entry holds. It reads
228# what the format writes; it is not a JSON validator: an escaped character is kept as itself (\uXXXX is not decoded),
229# and only strings, whole numbers, objects and lists are understood (anything else marks the header malformed).
230# =====================================================================================================
231
232# one tensor the model expects, in Layer.dump order: its name in the file and its shape (rows, columns, ...)
233type Entry is Data:
234  Entry{name: String, shape: List<&2, U32>}
235
236def St.entry_name(e: Entry) -> String:
237  match e:
238    case Entry{n, sh}: n
239
240type Tok is Data:
241  TStr{s: String}
242  TNum{v: W64}
243  TPun{c: U32}
244  TBad{c: U32}
245
246# the lexer's state between bytes: outside a token, inside a string (its characters reversed), just after a backslash
247# in a string, inside a number
248type Lex is Data:
249  LOut{}
250  LStr{acc: String}
251  LEsc{acc: String}
252  LNum{v: W64}
253
254# a state and the tokens one byte completed, in order
255type LexR is Data:
256  LexR{st: Lex, toks: List<&2, Tok>}
257
258def St.digit(+b: U32) -> Bool:
259  U32.is_lt(U32.sub(b, 48), 10)
260
261def St.space(+b: U32) -> Bool:
262  Bool.or(Bool.or(U32.is_eq(b, 32), U32.is_eq(b, 10)), Bool.or(U32.is_eq(b, 13), U32.is_eq(b, 9)))
263
264# { } [ ] : ,
265def St.pun(+b: U32) -> Bool:
266  Bool.or(Bool.or(Bool.or(U32.is_eq(b, 123), U32.is_eq(b, 125)), Bool.or(U32.is_eq(b, 91), U32.is_eq(b, 93))), Bool.or(U32.is_eq(b, 58), U32.is_eq(b, 44)))
267
268def St.out_go(quote: Bool, digit: Bool, space: Bool, pun: Bool, +b: U32) -> LexR:
269  match quote digit space pun:
270    case True{} _ _ _: LexR{LStr{""}, Nil{}}
271    case _ True{} _ _: LexR{LNum{W64.of(U32.sub(b, 48))}, Nil{}}
272    case _ _ True{} _: LexR{LOut{}, Nil{}}
273    case _ _ _ True{}: LexR{LOut{}, Con{TPun{b}, Nil{}}}
274    case _ _ _ _: LexR{LOut{}, Con{TBad{b}, Nil{}}}
275
276def St.out(+b: U32) -> LexR:
277  St.out_go(U32.is_eq(b, 34), St.digit(b), St.space(b), St.pun(b), b)
278
279def St.str_go(quote: Bool, slash: Bool, +b: U32, acc: String) -> LexR:
280  match quote slash:
281    case True{} _: LexR{LOut{}, Con{TStr{String.reverse(acc)}, Nil{}}}
282    case _ True{}: LexR{LEsc{acc}, Nil{}}
283    case _ _: LexR{LStr{SCon{Chr{b}, acc}}, Nil{}}
284
285# a number ends at the first byte that is not a digit, and that byte is then read as outside a token
286def St.num_end(+v: W64, r: LexR) -> LexR:
287  match r:
288    case LexR{st, ts}: LexR{st, Con{TNum{v}, ts}}
289
290def St.num_go(digit: Bool, +b: U32, +v: W64) -> LexR:
291  match digit:
292    case True{}: LexR{LNum{W64.add(W64.mul_small(v, 10), W64.of(U32.sub(b, 48)))}, Nil{}}
293    case False{}: St.num_end(v, St.out(b))
294
295def St.step(st: Lex, +b: U32) -> LexR:
296  match st:
297    case LOut{}: St.out(b)
298    case LStr{acc}: St.str_go(U32.is_eq(b, 34), U32.is_eq(b, 92), b, acc)
299    case LEsc{acc}: LexR{LStr{SCon{Chr{b}, acc}}, Nil{}}
300    case LNum{+v}: St.num_go(St.digit(b), b, v)
301
302# ts onto out, which holds the tokens so far in reverse
303def St.push(ts: List<&2, Tok>, out: List<&2, Tok>) -> List<&2, Tok>:
304  match ts:
305    case Nil{}: out
306    case Con{t, r}: St.push(r, Con{t, out})
307
308# the end of the header: a number ends there, an unfinished string is malformed
309def St.fin(st: Lex, out: List<&2, Tok>) -> List<&2, Tok>:
310  match st:
311    case LOut{}: out
312    case LNum{v}: Con{TNum{v}, out}
313    case LStr{acc}: Con{TBad{0}, out}
314    case LEsc{acc}: Con{TBad{0}, out}
315
316# the header's tokens, in order; r is what the previous byte left
317def St.lex(bs: List<&2, U32>, r: LexR, out: List<&2, Tok>) -> List<&2, Tok>:
318  match bs r:
319    case Nil{} LexR{st, ts}: List.reverse(&2, Tok, St.fin(st, St.push(ts, out)))
320    case Con{+b, t} LexR{st, ts}: St.lex(t, St.step(st, b), St.push(ts, out))
321
322# one top-level entry of the header: a tensor (or __metadata__, which no layout names)
323type Rec is Data:
324  Rec{name: String, dtype: String, shape: List<&2, W64>, offs: List<&2, W64>}
325
326# the pass over the tokens: the depth ({ and [ open a level), whether the last token was a colon, the last string,
327# the current entry's name and field, what the entry holds so far, the numbers of the open list, the entries read
328# (last first), and the count of tokens the header may not hold
329type Ps is Data:
330  Ps{depth: U32, colon: Bool, pend: String, name: String, field: String, dtype: String, shape: List<&2, W64>, offs: List<&2, W64>, nums: List<&2, W64>, recs: List<&2, Rec>, bad: U32}
331
332# a string is a value after a colon (the dtype, at depth 2) or else a key, kept until its colon
333def Sp.str(colon: Bool, is_dt: Bool, ps: Ps, s: String) -> Ps:
334  match colon is_dt ps:
335    case True{} True{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, s, sh, of, ns, rs, bad}
336    case True{} False{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, dt, sh, of, ns, rs, bad}
337    case False{} _ Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, s, n, f, dt, sh, of, ns, rs, bad}
338
339# a colon makes the last string the entry's name (depth 1) or the field's (depth 2)
340def Sp.colon(d1: Bool, d2: Bool, ps: Ps) -> Ps:
341  match d1 d2 ps:
342    case True{} _ Ps{d, c, +p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, True{}, p, p, f, dt, sh, of, ns, rs, bad}
343    case _ True{} Ps{d, c, +p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, True{}, p, n, p, dt, sh, of, ns, rs, bad}
344    case _ _ Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, True{}, p, n, f, dt, sh, of, ns, rs, bad}
345
346# { opens a level; at depth 2 an entry starts empty
347def Sp.open(fresh: Bool, ps: Ps) -> Ps:
348  match fresh ps:
349    case True{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{(d + 1 : U32), False{}, p, n, f, "", Nil{}, Nil{}, ns, rs, bad}
350    case False{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{(d + 1 : U32), False{}, p, n, f, dt, sh, of, ns, rs, bad}
351
352# } closes a level; closing depth 2 ends an entry
353def Sp.close(emit: Bool, ps: Ps) -> Ps:
354  match emit ps:
355    case True{} Ps{d, c, p, +n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, "", Nil{}, Nil{}, ns, Con{Rec{n, dt, sh, of}, rs}, bad}
356    case False{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, dt, sh, of, ns, rs, bad}
357
358def Sp.lopen(ps: Ps) -> Ps:
359  match ps:
360    case Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{(d + 1 : U32), False{}, p, n, f, dt, sh, of, Nil{}, rs, bad}
361
362# ] closes a list; at depth 3 it is the entry's shape or data_offsets
363def Sp.lclose(is_sh: Bool, is_of: Bool, ps: Ps) -> Ps:
364  match is_sh is_of ps:
365    case True{} _ Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, dt, List.reverse(&2, W64, ns), of, Nil{}, rs, bad}
366    case _ True{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, dt, sh, List.reverse(&2, W64, ns), Nil{}, rs, bad}
367    case _ _ Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{U32.sub(d, 1), False{}, p, n, f, dt, sh, of, Nil{}, rs, bad}
368
369def Sp.comma(ps: Ps) -> Ps:
370  match ps:
371    case Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, dt, sh, of, ns, rs, bad}
372
373# a number counts inside a list at depth 3; elsewhere it is not part of a tensor's entry
374def Sp.num(in_list: Bool, v: W64, ps: Ps) -> Ps:
375  match in_list ps:
376    case True{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, dt, sh, of, Con{v, ns}, rs, bad}
377    case False{} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}: Ps{d, False{}, p, n, f, dt, sh, of, ns, rs, bad}
378
379def Sp.pun(colon: Bool, ob: Bool, cb: Bool, ol: Bool, cl: Bool, +d: U32, +f: String, ps: Ps) -> Ps:
380  match colon ob cb ol cl:
381    case True{} _ _ _ _: Sp.colon(U32.is_eq(d, 1), U32.is_eq(d, 2), ps)
382    case _ True{} _ _ _: Sp.open(U32.is_eq(d, 1), ps)
383    case _ _ True{} _ _: Sp.close(U32.is_eq(d, 2), ps)
384    case _ _ _ True{} _: Sp.lopen(ps)
385    case _ _ _ _ True{}: Sp.lclose(Bool.and(U32.is_eq(d, 3), String.eq(f, "shape")), Bool.and(U32.is_eq(d, 3), String.eq(f, "data_offsets")), ps)
386    case _ _ _ _ _: Sp.comma(ps)
387
388def Sp.step(t: Tok, ps: Ps) -> Ps:
389  match t ps:
390    case TStr{s} Ps{+d, +c, p, n, +f, dt, sh, of, ns, rs, bad}:
391      Sp.str(c, Bool.and(U32.is_eq(d, 2), String.eq(f, "dtype")), Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}, s)
392    case TNum{v} Ps{+d, c, p, n, f, dt, sh, of, ns, rs, bad}:
393      Sp.num(U32.is_eq(d, 3), v, Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad})
394    case TPun{+k} Ps{+d, c, p, n, +f, dt, sh, of, ns, rs, bad}:
395      Sp.pun(U32.is_eq(k, 58), U32.is_eq(k, 123), U32.is_eq(k, 125), U32.is_eq(k, 91), U32.is_eq(k, 93), d, f, Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad})
396    case TBad{k} Ps{d, c, p, n, f, dt, sh, of, ns, rs, bad}:
397      Ps{d, c, p, n, f, dt, sh, of, ns, rs, (bad + 1 : U32)}
398
399def Sp.parse(ts: List<&2, Tok>, ps: Ps) -> Ps:
400  match ts:
401    case Nil{}: ps
402    case Con{t, r}: Sp.parse(r, Sp.step(t, ps))
403
404def St.rec_name(r: Rec) -> String:
405  match r:
406    case Rec{n, dt, sh, of}: n
407
408def St.find_put(r: Rec, rest: Maybe<&2, Rec>, hit: Bool) -> Maybe<&2, Rec>:
409  match hit:
410    case True{}: Some{r}
411    case False{}: rest
412
413# the entry named `name`
414def St.find(rs: List<&2, Rec>, +name: String) -> Maybe<&2, Rec>:
415  match rs:
416    case Nil{}: None{}
417    case Con{+r, t}: St.find_put(r, St.find(t, name), String.eq(St.rec_name(r), name))
418
419def St.prod(xs: List<&2, U32>) -> U32:
420  match xs:
421    case Nil{}: 1
422    case Con{x, t}: (x * St.prod(t) : U32)
423
424def St.total(es: List<&2, Entry>) -> U32:
425  match es:
426    case Nil{}: 0
427    case Con{Entry{n, sh}, t}: (St.prod(sh) + St.total(t) : U32)
428
429def St.shape_eq(a: List<&2, W64>, b: List<&2, U32>) -> Bool:
430  match a b:
431    case Nil{} Nil{}: True{}
432    case Con{x, s} Con{y, t}: Bool.and(W64.eq(x, W64.of(y)), St.shape_eq(s, t))
433    case _ _: False{}
434
435def St.show_u(xs: List<&2, U32>, sep: String) -> String:
436  match xs:
437    case Nil{}: ""
438    case Con{x, t}: sep ++ U32.show(x) ++ St.show_u(t, ", ")
439
440def St.show_w(xs: List<&2, W64>, sep: String) -> String:
441  match xs:
442    case Nil{}: ""
443    case Con{x, t}: sep ++ W64.show(x) ++ St.show_w(t, ", ")
444
445# where a tensor is and how to read it, or why it cannot be read
446type Plan is Data:
447  PlanOk{name: String, bf16: Bool, off: W64, n: U32}
448  PlanBad{msg: String}
449
450def St.plan_span(ok: Bool, +name: String, bf16: Bool, off: W64, +n: U32) -> Plan:
451  match ok:
452    case True{}: PlanOk{name, bf16, off, n}
453    case False{}: PlanBad{"tensor " ++ name ++ ": data_offsets do not span its " ++ U32.show(n) ++ " numbers"}
454
455# data_offsets must be [begin, end] with end - begin the tensor's bytes
456def St.plan_of(+name: String, bf16: Bool, +start: W64, +n: U32, +width: U32, of: List<&2, W64>) -> Plan:
457  match of:
458    case Con{+b, Con{e, Nil{}}}:
459      St.plan_span(W64.eq(W64.add(b, W64.mul_small(W64.of(n), width)), e), name, bf16, W64.add(start, b), n)
460    case _:
461      PlanBad{"tensor " ++ name ++ ": data_offsets is not [begin, end]"}
462
463def St.plan_shape(ok: Bool, +name: String, f32: Bool, +start: W64, +sh: List<&2, W64>, +want: List<&2, U32>, of: List<&2, W64>) -> Plan:
464  match ok f32:
465    case True{} True{}: St.plan_of(name, False{}, start, St.prod(want), 4, of)
466    case True{} False{}: St.plan_of(name, True{}, start, St.prod(want), 2, of)
467    case False{} _: PlanBad{"tensor " ++ name ++ ": shape [" ++ St.show_w(sh, "") ++ "] in the file, [" ++ St.show_u(want, "") ++ "] expected"}
468
469def St.plan_dtype(f32: Bool, bf16: Bool, +name: String, +dt: String, +start: W64, +sh: List<&2, W64>, +want: List<&2, U32>, of: List<&2, W64>) -> Plan:
470  match f32 bf16:
471    case True{} _: St.plan_shape(St.shape_eq(sh, want), name, True{}, start, sh, want, of)
472    case _ True{}: St.plan_shape(St.shape_eq(sh, want), name, False{}, start, sh, want, of)
473    case _ _: PlanBad{"tensor " ++ name ++ ": dtype " ++ dt ++ " (only F32 and BF16 are read)"}
474
475def St.plan(+start: W64, e: Entry, mr: Maybe<&2, Rec>) -> Plan:
476  match e mr:
477    case Entry{+name, want} None{}: PlanBad{"tensor " ++ name ++ ": not in the file"}
478    case Entry{+name, want} Some{Rec{rn, +dt, sh, of}}: St.plan_dtype(String.eq(dt, "F32"), String.eq(dt, "BF16"), name, dt, start, sh, want, of)
479
480def St.read(bf16: Bool, f: File, off: W64, +n: U32, a: Array<F32>, +at: U32) -> IO(File & (Array<F32> & Result<&1, &1, U32 & String, U32>)):
481  match bf16 off:
482    case True{} W64{hi, lo}: R.read_bf16_into(f, hi, lo, n, a, at)
483    case False{} W64{hi, lo}: R.read_f32_into(f, hi, lo, n, a, at)
484
485def St.got(+path: String, +name: String, +n: U32, r: File & (Array<F32> & Result<&1, &1, U32 & String, U32>), k: File -> Array<F32> -> U32 -> IO(Array<F32>)) -> IO(Array<F32>):
486  match r:
487    case Tuple{f, Tuple{a, res}}:
488      do IO<Array<F32>>:
489        +got : U32 <- IO.pass(U32, res)
490        require(U32.is_eq(got, n), path ++ ": tensor " ++ name ++ ": short read (" ++ U32.show(got) ++ " of " ++ U32.show(n) ++ " numbers)")
491        k(f, a, n)
492
493# reads one planned tensor into cells at .., then goes on with k
494def St.then(+path: String, p: Plan, f: File, a: Array<F32>, +at: U32, k: File -> Array<F32> -> U32 -> IO(Array<F32>)) -> IO(Array<F32>):
495  match p:
496    case PlanBad{msg}: IO.die(Array<F32>, 1, path ++ ": " ++ msg)
497    case PlanOk{+name, bf16, off, +n}:
498      IO.bind(File & (Array<F32> & Result<&1, &1, U32 & String, U32>), Array<F32>, St.read(bf16, f, off, n, a, at), r => St.got(path, name, n, r, k))
499
500# every tensor of es, in order, into consecutive cells from at
501def St.fill(es: List<&2, Entry>, +path: String, +start: W64, +recs: List<&2, Rec>, f: File, a: Array<F32>, +at: U32) -> IO(Array<F32>):
502  match es:
503    case Nil{}:
504      do IO<Array<F32>>:
505        File.close(f)
506        return a
507    case Con{+e, t}:
508      St.then(path, St.plan(start, e, St.find(recs, St.entry_name(e))), f, a, at, f2 => a2 => +n => St.fill(t, path, start, recs, f2, a2, (at + n : U32)))
509
510def St.le32(+a: U32, +b: U32, +c: U32, +d: U32) -> U32:
511  U32.or(U32.or(a, U32.shln(b, 8n)), U32.or(U32.shln(c, 16n), U32.shln(d, 24n)))
512
513# the header's length: the file's first 8 bytes, little-endian
514def St.le64(bs: List<&2, U32>) -> Maybe<&2, W64>:
515  match bs:
516    case Con{b0, Con{b1, Con{b2, Con{b3, Con{b4, Con{b5, Con{b6, Con{b7, t}}}}}}}}:
517      Some{W64{St.le32(b4, b5, b6, b7), St.le32(b0, b1, b2, b3)}}
518    case _:
519      None{}
520
521def St.count(bs: List<&2, U32>, +acc: U32) -> U32:
522  match bs:
523    case Nil{}: acc
524    case Con{b, t}: St.count(t, (acc + 1 : U32))
525
526# A parsed safetensors header: where the data starts and each tensor's record. Reading it once lets a model be filled
527# one piece at a time (`piece`), so that no array of the whole model exists. One array for all of SmolLM2-1.7B's
528# 1.71e9 parameters takes a 2^31-cell block (8 GiB) whose memory the training step's smaller blocks never reuse.
529type Header is Data:
530  Header{start: W64, recs: List<&2, Rec>}
531
532def St.parsed(+path: String, +start: W64, ps: Ps) -> IO(Header):
533  match ps:
534    case Ps{d, c, p, nm, fl, dt, sh, of, ns, rs, bad}:
535      do IO<Header>:
536        require(Bool.and(U32.is_zero(d), U32.is_zero(bad)), path ++ ": malformed safetensors header")
537        return Header{start, rs}
538
539def St.header(+path: String, +len: U32, r: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(Header):
540  match r:
541    case Tuple{f, res}:
542      do IO<Header>:
543        +bs : List<&2, U32> <- IO.pass(List<&2, U32>, res)
544        File.close(f)
545        require(U32.is_eq(St.count(bs, 0), len), path ++ ": the header is cut short")
546        St.parsed(path, W64.add(W64.of(8), W64.of(len)), Sp.parse(St.lex(bs, LexR{LOut{}, Nil{}}, Nil{}), Ps{0, False{}, "", "", "", "", Nil{}, Nil{}, Nil{}, Nil{}, 0}))
547
548# the format caps the header at 100 MB
549def St.length(f: File, +path: String, h: Maybe<&2, W64>) -> IO(Header):
550  match h:
551    case None{}: IO.die(Header, 1, path ++ ": shorter than a safetensors header")
552    case Some{W64{hi, +lo}}:
553      do IO<Header>:
554        require(Bool.and(U32.is_zero(hi), U32.is_le(lo, 100000000)), path ++ ": a safetensors header past 100 MB")
555        r : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, 8, lo)
556        St.header(path, lo, r)
557
558def St.first8(+path: String, r: File & Result<&1, &1, U32 & String, List<&2, U32>>) -> IO(Header):
559  match r:
560    case Tuple{f, res}:
561      do IO<Header>:
562        bs : List<&2, U32> <- IO.pass(List<&2, U32>, res)
563        St.length(f, path, St.le64(bs))
564
565# A safetensors file's header, parsed; the file is closed after it
566def header(+path: String) -> IO(Header):
567  do IO<Header>:
568    f : File <- IO.try(File, File.open(path, "r"))
569    r : File & Result<&1, &1, U32 & String, List<&2, U32>> <- File.read_at(f, 0, 8)
570    St.first8(path, r)
571
572# The tensors `layout` names, read from the file whose header is h, in the layout's order, into one array of n numbers
573def tensors_from(+path: String, h: Header, +n: U32, +layout: List<&2, Entry>) -> IO(Array<F32>):
574  match h:
575    case Header{start, recs}:
576      do IO<Array<F32>>:
577        require(U32.is_eq(St.total(layout), n), path ++ ": the layout holds " ++ U32.show(St.total(layout)) ++ " numbers, the model " ++ U32.show(n))
578        f : File <- IO.try(File, File.open(path, "r"))
579        St.fill(layout, path, start, recs, f, Array.new(F32, T.depth(U32.to_nat(n)), 0.0), 0)
580
581# The tensors `layout` names, read from a safetensors file in its order into one array of n numbers (F32 as they are,
582# BF16 widened exactly). Each must be in the file with the shape given, as F32 or BF16, and the shapes must hold n
583# numbers in all; otherwise the program stops with a message naming the tensor. Other tensors in the file are ignored.
584def tensors(+path: String, +n: U32, +layout: List<&2, Entry>) -> IO(Array<F32>):
585  do IO<Array<F32>>:
586    require(U32.is_eq(St.total(layout), n), path ++ ": the layout holds " ++ U32.show(St.total(layout)) ++ " numbers, the model " ++ U32.show(n))
587    h : Header <- header(path)
588    tensors_from(path, h, n, layout)
589
590# One layer's parameters from the file whose header is h: `layout` names its tensors in Layer.dump order, and
591# Layer.load fills them; the array they were read into is gone once this answers
592def piece(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer<P, X, Y, S>, +path: String, h: Header, +layout: List<&2, Entry>, p: P) -> IO(P):
593  do IO<P>:
594    +n : U32 = T.Layer.count(P, X, Y, S, m)
595    a : Array<F32> <- tensors_from(path, h, n, layout)
596    weights_done(P, n, T.Layer.load(P, X, Y, S, m, p, a, 0))
597
598# A pair of parameters filled one after the other: the first piece's source array is released before the second is read
599def two(-P1: Type, -P2: Type, p: T.Two<P1, P2>, k1: P1 -> IO(P1), k2: P2 -> IO(P2)) -> IO(T.Two<P1, P2>):
600  match p:
601    case T.Two{a, b}:
602      do IO<T.Two<P1, P2>>:
603        x : P1 <- k1(a)
604        y : P2 <- k2(b)
605        return T.Two{x, y}
606
607# All parameters from a safetensors file: `layout` names the model's tensors in Layer.dump order (a tied output head
608# is not in it, as SmolLM2's is not in its file), and Layer.load fills the parameters from them as `weights` does
609def safetensors(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer<P, X, Y, S>, +path: String, +layout: List<&2, Entry>, p: P) -> IO(P):
610  do IO<P>:
611    h : Header <- header(path)
612    piece(~P, ~X, ~Y, ~S, ~m, path, h, layout, p)