Tensor-Core Fragment Shape

Problem

Compute the work of one tensor-core matrix-multiply instruction of shape m x n x k: macs = m * n * k multiply-accumulates and flops = 2 * macs.

The signature is tensor_core_fragment(m, n, k). Return [macs, flops].

Hint

One multiply-accumulate counts as two floating point operations.

case 1

in: [16,8,16]

expected: [2048,4096]

case 2

in: [16,16,16]

expected: [4096,8192]

case 3

in: [8,8,4]

expected: [256,512]

case 4

in: [1,1,1]

expected: [1,2]

Starter codePython

def tensor_core_fragment(m, n, k):
    # Your code here
    pass

4 test cases · run your code against them in the editor