In Tensorflow, how to use tf.gather() for the last dimension?

Question

I am trying to gather slices of a tensor in terms of the last dimension for partial connection between layers. Because the output tensor's shape is [batch_size, h, w, depth], I want to select slices based on the last dimension, such as

# L is intermediate tensor
partL = L[:, :, :, [0,2,3,8]]

However, tf.gather(L, [0, 2,3,8]) seems to only work for the first dimension (right?) Can anyone tell me how to do it?

Answer 1

As of TensorFlow 1.3 tf.gather has an axis parameter, so the various workarounds here are no longer necessary.

https://www.tensorflow.org/versions/r1.3/api_docs/python/tf/gather https://github.com/tensorflow/tensorflow/issues/11223

Answer 2

There's a tracking bug to support this use-case here: https://github.com/tensorflow/tensorflow/issues/206

For now you can:

1. transpose your matrix so that dimension to gather is first (transpose is expensive)

2. reshape your tensor into 1d (reshape is cheap) and turn your gather column indices into a list of individual element indices at linear indexing, then reshape back

3. use gather_nd. Will still need to turn your column indices into list of individual element indices.

Answer 3

With gather_nd you can now do this as follows:

cat_idx = tf.concat([tf.range(0, tf.shape(x)[0]), indices_for_dim1], axis=0)
result = tf.gather_nd(matrix, cat_idx)

Also, as reported by user Nova in a thread referenced by @Yaroslav Bulatov's:

x = tf.constant([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9]])
idx = tf.constant([1, 0, 2])
idx_flattened = tf.range(0, x.shape[0]) * x.shape[1] + idx
y = tf.gather(tf.reshape(x, [-1]),  # flatten input
              idx_flattened)  # use flattened indices

with tf.Session(''):
  print y.eval()  # [2 4 9]

The gist is flatten the tensor and use strided 1D addressing with tf.gather(...).