~/jendhub

def attn_fwd

bendtorch@1.3.0.0 · main.bend:2945

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

2945def attn_fwd(+n: Nat, +w: Nat, +kvw: Nat, +dm: AttnDims, x: Two<TA.Mat<n, w>, Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>) -> Fwd<Unit, TA.Mat<n, w>, AttnS>:
2946  match dm x:
2947    case AttnDims{+b, +t, +kv, +g, +hd, +sc} Two{TA.Mat{q}, Two{TA.Mat{k}, TA.Mat{v}}}:
2948      af2(n, w, b, kv, dm_of(dm), fslices(b, 0n, t, w, kvw, kv, g, hd, FL{q, k, v, Nil{}}))

Uses (12)

Used by (1)