~/jendhub

def ab2

bendtorch@1.3.0.0 · main.bend:3004

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

3004def ab2(+n: Nat, +w: Nat, +kvw: Nat, +b: Nat, +kv: Nat, +dm: Dm, s: BL) -> Bwd<Unit, Two<TA.Mat<n, w>, Two<TA.Mat<n, kvw>, TA.Mat<n, kvw>>>>:
3005  match dm s:
3006    case Dm{+t, +g, +hd, +gw, +kw, +sc} BL{q, k, v, d, xs}:
3007      ab3(n, w, kvw, dblocks_in(b, 0n, t, gw, hd, w, kvw, kv, UD{bgo(attn_tree(Nat.mul(b, kv), t, gw), Nat.mul(b, kv), 1n, dm, List.reverse(&1, BG, xs)), D3{zeros_a(Nat.mul(n, w)), zeros_a(Nat.mul(n, kvw)), zeros_a(Nat.mul(n, kvw))}}))

Uses (17)

Used by (1)