Pruning Mask Apply

Problem

Apply a pruning mask to a weight matrix element-wise: pruned_weights[i][j] = weights[i][j] * mask[i][j].

The signature is apply_pruning_mask(weights, mask). Both matrices have the same shape.

Hint

A mask entry of 1 keeps the weight and 0 zeroes it out.

case 1

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

expected: [[1,0],[0,4]]

case 2

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

expected: [[0,0]]

case 3

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

expected: [[-1,2],[3,-4]]

case 4

in: [[[2,3],[-5,6]],[[0,1],[0,0]]]

expected: [[0,3],[0,0]]

Starter codePython

def apply_pruning_mask(weights, mask):
    # Your code here
    pass

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