I have been using a tensor as a mask to select specific values from a matrix.
Below an example using NumPy and converting to Tensorflow.
a = np.zeros((
