def af2_full
bendtorch@1.3.0.0 · main.bend:2989
import bendtorch@1.3.0.0/main.bend as Bendtorch
import Base
import bend-ml-tensor-array@0.1.18.0/main.bend as TA
2989def af2_full(-n: Nat, +w: Nat, +b: Nat, +kv: Nat, +dm: Dm, s: FL) -> Fwd<Unit, TA.Mat<n, w>, AttnS>:
2990 match dm s:
2991 case Dm{+t, +g, +hd, +gw, +kw, +sc} FL{q, k, v, xs}:
2992 Fwd{Unit{}, TA.Mat{blocks_in(b, 0n, t, gw, w, kv, UO{fgo_full(attn_tree(Nat.mul(b, kv), t, gw), Nat.mul(b, kv), g, dm, List.reverse(&1, FS, xs)), zeros_a(Nat.mul(Nat.mul(b, t), w))})}, AttnS{q, k, v}}