def attn_bwd_full

bendtorch@1.3.0.0 · main.bend:3083

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

3083def attn_bwd_full(+n: Nat, +w: Nat, +kvw: Nat, +dm: AttnDims, s: AttnS, dy: TA.Mat<n, w>) -> Bwd<Unit, Two<TA.Mat<n, w>, Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
3084  match dm s dy:
3085    case AttnDims{+b, +t, +kv, +g, +hd, +sc} AttnS{q, k, v} TA.Mat{d}:
3086      ab2_full(n, w, kvw, b, kv, dm_of(dm), bslices(b, 0n, t, w, kvw, kv, g, hd, BL{q, k, v, d, Nil{}}))

Uses (12)

Used by (1)