Fancy indexing in Tensorflow: Getting a different element from every row in a matrix

Let's say you have the following 4 by 3 matrix:

M = np.array([[  1,  2,  3 ],
              [  4,  5,  6 ],
              [  7,  8,  9 ],
              [ 10, 11, 12 ]])

When in Tensorflow we'd also have this line:

M = tf.constant(M)

Let's say that you want to get the first element of every row. In both Numpy and Tensorflow you can just do this:

x = M[:, 0]

which means 'get every row and from every row get the element at index 0'. "x" is now equal to:

np.array([1, 4, 7, 10])

Now let's say that instead of the first element of every row, you wanted the third element of the first row, the second element of the second row, the first element of the third row, and the second element of the fourth row. In Numpy, this is how you do that:

idx = [2,1,0,1]
x = M[[0,1,2,3], idx]

or equivalently:

x = M[np.arange(M.shape[0]), idx]

This is just breaking up the 'coordinates' of the desired elements into separate lists for each axis. "x" is now equal to:

np.array([3, 5, 7, 11])

Unfortunately this sort of fancy indexing isn't available in Tensorflow. Instead we have the function "tf.gather_nd" which lets us provide a list of 'coordinates'. Unlike Numpy, tf.gather_nd does not take separate lists per axis but expects a single list with one coordinate per item like this:

idx = tf.constant([[0,2],[1,1],[2,0],[3,1]], tf.int32)
x = tf.gather_nd(M, idx)

This is typically inconvenient as we usually have a single vector of indexes rather then a list of coordinates. It would be better to be able to just put a "range" like we did with Numpy. We can use the range and then join it to the vector of indexes sideways using "tf.stack", like this:

idx = tf.constant([2,1,0,1], tf.int32)
x = tf.gather_nd(M, tf.stack([tf.range(M.shape[0]), idx], axis=1))

Kind of bulky but at least it's possible. I miss Theano's Numpy-like interface.