def matmul_tn

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

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

264# C = Aᵀ · B   (k x n)ᵀ · (k x m): a weight gradient, Xᵀ · dY
265def matmul_tn(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<k, n>, b: TA.Mat<k, m>) -> TA.MMulTN<n, k, m>:
266  match a b:
267    case TA.Mat{x} TA.Mat{y}:
268      TA.mmul_tn_fin(n, k, m, gemm_any(False{}, TA.split_depth(par, n, m, k), n, m, k, 1, U32.from_nat(n), U32.from_nat(m), 1, x, y, Array.new(F32, depth(Nat.mul(n, m)), 0.0)))