def sm_loop

bendtorch-gpu@1.3.0.0 · main.bend:783

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

783# forward 2: P's row r (query i = r % t) by main.bend's softmax over its i + 1 positions
784def sm_loop(k: Nat, +r: U32, +l: U32, +t32: U32, +tr: U32, p: Array<F32>) -> Array<F32>:
785  match k:
786    case 0n: p
787    case 1n+q: sm_loop(q, (r + l : U32), l, t32, tr, T.pc_a(T.psmc(U32.to_nat((U32.mod(r, t32) + 1 : U32)), toff(U32.div(r, t32), U32.mod(r, t32), tr), p)))