def prow
import bendtorch@1.3.0.0/main.bend as Bendtorchimport Base
2426# row i of a head's P (t x t, row-major, created zeros): the scaled scores of query i against keys 0..i (qrow) written
2427# from the row's start, then the softmax over as many positions as were written: i + 1 (the law causal). The positions
2428# after i are never written and keep their zeros.
2429def prow(+i: Nat, +t: Nat, +hd: Nat, +sc: F32, +qs: U32, +qc: U32, +ks: U32, +kc: U32, q: Array<F32>, k: Array<F32>, p: Array<F32>) -> PR:
2430 prow2(sc, (U32.from_nat(i) * U32.from_nat(t) : U32), p, qrow(i, hd, qs, qc, ks, kc, q, k))