main.bend
bendtorch-load@1.0.0.0 · 68 lines · 0 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).
3import Base
4import bend-ml-binary-read@0.1.0.0/main.bend as R
5import bendtorch@1.0.0.0/main.bend as T
6
7def require(ok: Bool, +message: String) -> IO(Unit):
8 match ok:
9 case True{}: IO.pure(Unit, Unit{})
10 case False{}: IO.die(Unit, 1, message)
11
12def read_done(+path: String, +n: U32, r: File & (Array<F32> & Result<&1, &1, U32 & String, U32>)) -> IO(Array<F32>):
13 match r:
14 case Tuple{f, Tuple{a, res}}:
15 do IO<Array<F32>>:
16 +got : U32 <- IO.pass(U32, res)
17 File.close(f)
18 require(U32.is_eq(got, (4 * n : U32)), path ++ ": short F32 read")
19 return a
20
21def read_sized(+path: String, +n: U32, r: File & Result<&1, &1, U32 & String, U32>) -> IO(Array<F32>):
22 match r:
23 case Tuple{f, res}:
24 do IO<Array<F32>>:
25 +size : U32 <- IO.pass(U32, res)
26 require(U32.is_eq(size, (4 * n : U32)), path ++ ": expected " ++ U32.show((4 * n : U32)) ++ " bytes, got " ++ U32.show(size))
27 rr : File & (Array<F32> & Result<&1, &1, U32 & String, U32>) <- R.read_f32(f, 0, n, Array.new(F32, T.depth(U32.to_nat(n)), 0.0))
28 read_done(path, n, rr)
29
30# A raw, exact-size little-endian F32 file, read once into one source array. Counts and file offsets are U32.
31def read(+path: String, +n: U32) -> IO(Array<F32>):
32 do IO<Array<F32>>:
33 require(U32.is_le(n, 1073741823), "F32 file exceeds the binary reader's U32 byte limit")
34 f : File <- IO.try(File, File.open(path, "r"))
35 rr : File & Result<&1, &1, U32 & String, U32> <- File.size(f)
36 read_sized(path, n, rr)
37
38def weights_done(-P: Type, +n: U32, r: T.PA<P>) -> IO(P):
39 match r:
40 case T.PA{p, xs, +off}:
41 do IO<P>:
42 require(U32.is_eq(off, n), "Layer.count differs from the loading cursor")
43 return p
44
45# All parameters in Layer.dump order. The source array is released after the model's arrays are filled.
46def weights(~P: Type, ~X: Type, ~Y: Type, ~S: Type, ~m: T.Layer<P, X, Y, S>, +path: String, p: P) -> IO(P):
47 do IO<P>:
48 +n : U32 = T.Layer.count(P, X, Y, S, m)
49 a : Array<F32> <- read(path, n)
50 weights_done(P, n, T.Layer.load(P, X, Y, S, m, p, a, 0))
51
52# The checkpoint layout is determined by the already initialized optimizer state, without a header:
53# SGD: parameters; momentum: parameters, buffer; Adam/AdamW: parameters, m, v, 1 - b1^t, 1 - b2^t.
54def 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>):
55 do IO<T.Train<P>>:
56 require(U32.is_le(base, U32.div(U32.sub(1073741823, tail), copies)), "checkpoint exceeds the binary reader's U32 byte limit")
57 +n : U32 = ((copies * base : U32) + tail : U32)
58 a : Array<F32> <- read(path, n)
59 return T.Train.restore(~P, ~X, ~Y, ~S, ~m, tr, a)
60
61def 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>):
62 match tr:
63 case T.Train{p, g, T.NoMom{}, p1, p2, loss}:
64 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})
65 case T.Train{p, g, T.OneMom{mm}, p1, p2, loss}:
66 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})
67 case T.Train{p, g, T.TwoMom{mm, vv}, p1, p2, loss}:
68 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})