Tensor Contraction Cost

Problem

Count the scalar multiplications needed to contract two tensors: tensor 1 with the given shape contracts its dimension axis1 with dimension axis2 of tensor 2.

The cost is (product of all dimensions of shape1) * (product of the free dimensions of shape2).

Hint

Matrix multiplication is the special case m*k times k*n, costing m*k*n multiplications.

Test cases

case 1

in: [[2,3],[3,4],1,0]

expected: 24

case 2

in: [[5,2],[2,2],1,0]

expected: 20

case 3

in: [[2,2,2],[2,2],2,0]

expected: 16

case 4

in: [[1,1],[1,1],0,0]

expected: 1

Starter codePython

def contraction_cost(shape1, shape2, axis1, axis2):
    # Your code here
    pass

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