def matmul_acc

bend-ml-tensor-fork@0.1.6.0 · main.bend:246

import Base
import bend-ml-tensor-array@0.1.19.0/main.bend as TA

246# C + A · B, into C
247def matmul_acc(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<k, m>, c: TA.Mat<n, m>) -> TA.MMul<n, k, m>:
248  match a b c:
249    case TA.Mat{x} TA.Mat{y} TA.Mat{z}:
250      TA.mmul_fin(n, k, m, gemm_any(True{}, TA.split_depth(par, n, m, k), n, m, k, U32.from_nat(k), 1, U32.from_nat(m), 1, x, y, z))