def load_lw_st

bend-ml-llama@0.2.1.0 · main.bend:1666

import Base
import bend-ml-tensor-array@0.1.19.0/main.bend as TA

1666# one layer's weights (tensors t_layer(i, 0..8)) from w/st.txt's numbers st; d: the band depth (depth_for)
1667def load_lw_st(+c: Cfg, +dir: String, +st: List<&2, U32>, +i: Nat, +d: Nat) -> IO(LW<c>):
1668  do IO<LW<c>>:
1669    ln1 : List<&2, F32> <- get_norm(dir, st, t_layer(i, 0n), Cfg.dim(c))
1670    q : TA.Bands<Cfg.dim(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 1n), d, Cfg.dim(c), Cfg.dim(c))
1671    k : TA.Bands<Cfg.kvw(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 2n), d, Cfg.kvw(c), Cfg.dim(c))
1672    v : TA.Bands<Cfg.kvw(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 3n), d, Cfg.kvw(c), Cfg.dim(c))
1673    o : TA.Bands<Cfg.dim(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 4n), d, Cfg.dim(c), Cfg.dim(c))
1674    ln2 : List<&2, F32> <- get_norm(dir, st, t_layer(i, 5n), Cfg.dim(c))
1675    gate : TA.Bands<Cfg.ffn(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 6n), d, Cfg.ffn(c), Cfg.dim(c))
1676    up : TA.Bands<Cfg.ffn(c), Cfg.dim(c)> <- get_bands(dir, st, t_layer(i, 7n), d, Cfg.ffn(c), Cfg.dim(c))
1677    down : TA.Bands<Cfg.dim(c), Cfg.ffn(c)> <- get_bands(dir, st, t_layer(i, 8n), d, Cfg.dim(c), Cfg.ffn(c))
1678    return LW{ln1, q, k, v, o, ln2, gate, up, down}