def head_of

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

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

1018# the final norm, then logits = emb · nf (the LM head is the embedding matrix)
1019def head_of(+c: Cfg, -cap: Nat, -t: Nat, +d: Nat, emb: TA.Bands<Cfg.vocab(c), Cfg.dim(c)>, +lnf: List<&2, F32>, r: Run<c, cap, t>) -> Out<c, cap, t>:
1020  match r:
1021    case Run{+x, ls, cs, acts}:
1022      +nf = rmsnorm(c, x, lnf)
1023      mv(Out<c, cap, t>, Cfg.vocab(c), Cfg.dim(c), emb, nf, emb => +lg => Out{lg, nf, acts, cs, Model{d, emb, ls, lnf}})