def KV.push

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

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

499# position t: the keys and values of every key/value head written at t·head; room is the proof that t < cap
500def KV.push(+c: Cfg, -cap: Nat, +t: Nat, kv: KV<cap, t>, -room: {True{} == TA.leb(1n+t, cap) : Bool}, +ks: List<&2, List<&2, F32>>, +vs: List<&2, List<&2, F32>>) -> KV<cap, 1n+t>:
501  match kv:
502    case KV{ka, va}:
503      KV{write_rows(ka, ks, U32.from_nat(Nat.mul(t, Cfg.head(c)))), write_rows(va, vs, U32.from_nat(Nat.mul(t, Cfg.head(c))))}