def matmul_nt
import bend-ml-tensor-gpu@0.1.5.1/main.bend as BendMlTensorGpuimport Base
import bend-ml-tensor-array@0.1.19.0/main.bend as TA
304def matmul_nt(+n: Nat, +k: Nat, +m: Nat, par: Nat, a: TA.Mat<n, k>, b: TA.Mat<m, k>) -> TA.MMulNT<n, k, m>:
305 match a b:
306 case TA.Mat{x} TA.Mat{y}:
307 mm_nt_fin(n, k, m, gemm_new(n, m, k, U32.from_nat(k), 1, 1, U32.from_nat(k), x, y))