Mask and reshape a matrix using Tensorflow

Question

I would like to use a mask to select vectors in mat variable. Below is the resultant matrix that I expcxt to get:

[[0, 0, 2, 2],
 [4