BendMlLlama.head_of: def of bend-ml-llama@0.2.1.0
Import: import bend-ml-llama@0.2.1.0/main.bend as BendMlLlama
Lines 1018-1023 of bend-ml-llama@0.2.1.0/main.bend: https://jend.leeeet.dev/read.txt?path=bend-ml-llama@0.2.1.0/main.bend&start_line=1018&end_line=1023

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

# the final norm, then logits = emb · nf (the LM head is the embedding matrix)
def 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>:
  match r:
    case Run{+x, ls, cs, acts}:
      +nf = rmsnorm(c, x, lnf)
      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}})

Uses (12):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=List
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.dim
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Cfg.vocab
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Model
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Out
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=Run
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=mv
  https://jend.leeeet.dev/definition.txt?path=bend-ml-llama@0.2.1.0/main.bend&name=rmsnorm
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-array@0.1.19.0/main.bend&name=Bands