main.bend

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