def tsum_h

bendtorch-gpu@1.3.0.0 · main.bend:665

import Base
import bendtorch@1.3.0.0/main.bend as T

665# acc + the group's heads' sums, in head order (main.bend's pacc): a's column j (tr apart a head) · b's column (hd
666# apart a head), rows j..t-1
667def tsum_h(k: Nat, +ia: U32, +ib: U32, +j: U32, +tn: Nat, +w: U32, +sa: U32, +sb: U32, acc: F32, r: T.D2) -> T.D2:
668  match k r:
669    case 0n T.D2{a, b, +s}: T.D2{a, b, (acc + s : F32)}
670    case 1n+q T.D2{a, b, +s}: tsum_h(q, (ia + sa : U32), (ib + sb : U32), j, tn, w, sa, sb, (acc + s : F32), tdot(tn, (ia + sa : U32), (ib + sb : U32), j, w, 0.0, Array.get(F32, a, (ia + sa : U32)), Array.get(F32, b, (ib + sb : U32))))