def rope_go2

bendtorch-gpu@1.3.0.0 · main.bend:589

import Base
import bend-ml-tensor-array@0.1.18.0/main.bend as TA
import bend-ml-tensor-gpu@0.1.5.0/main.bend as TG
import bendtorch@1.3.0.0/main.bend as T

589def rope_go2(+n: Nat, +w: Nat, +fwd: Bool, +t: Nat, +h: Nat, +half: Nat, x: Array<F32>, tb: T.CS) -> TA.Mat<n, w>:
590  match tb:
591    case T.CS{cs, sn}: TA.Mat{r1_a(TG.run1(~ERO, ~lf_rope, TG.depth(Nat.mul(n, Nat.mul(h, half))), ERO{U32.from_nat(Nat.mul(n, Nat.mul(h, half))), fwd, U32.from_nat(Nat.mul(h, half)), U32.from_nat(half), U32.from_nat(w), U32.from_nat(t)}, cs, sn, nn(Unit{}), nn(Unit{}), x))}