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	import Base
4	import bend-ml-binary-read@0.1.0.0/main.bend as R
5	import bendtorch@1.0.1.0/main.bend as T
6	
7	def 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	
12	def 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	
21	def 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.
31	def 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	
38	def 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.
46	def 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.
54	def 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	
61	def 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})