tensorflow Tensor indexing Extract non-contiguous slices from the first dimension of a tensor

Example

Generally tf.gather gives you access to elements in the first dimension of a tensor (e.g. rows 1, 3 and 7 in a 2-dimensional Tensor). If you need access to any other dimension than the first one, or if you don't need the whole slice, but e.g. only the 5th entry in the 1st, 3rd and 7th row, you are better off using tf.gather_nd (see upcoming example for this).

tf.gather arguments:

• params: A tensor you want to extract values from.
• indices: A tensor specifying the indices pointing into params

Refer to the tf.gather(params, indices) documentation for detailed information.

---

We want to extract the 1st and 4th row in a 2-dimensional tensor.

# data is [[0, 1, 2, 3, 4, 5],
#          [6, 7, 8, 9, 10, 11],
#          ...
#          [24, 25, 26, 27, 28, 29]]
data = np.reshape(np.arange(30), [5, 6])
params = tf.constant(data)
indices = tf.constant([0, 3])
selected = tf.gather(params, indices)

selected has shape [2, 6] and printing its value gives

[[ 0  1  2  3  4  5]
 [18 19 20 21 22 23]]

---

indices can also just be a scalar (but cannot contain negative indices). E.g. in the above example:

tf.gather(params, tf.constant(3))

would print

[18 19 20 21 22 23]

---

Note that indices can have any shape, but the elements stored in indices always only refer to the first dimension of params. E.g. if you want to retrieve both the 1st and 3rd row and the 2nd and 4th row at the same time, you can do this:

indices = tf.constant([[0, 2], [1, 3]])
selected = tf.gather(params, indices)

Now selected will have shape [2, 2, 6] and its content reads:

[[[ 0  1  2  3  4  5]
  [12 13 14 15 16 17]]

 [[ 6  7  8  9 10 11]
  [18 19 20 21 22 23]]]

---

You can use tf.gather to compute a permutation. E.g. the following reverses all rows of params:

indices = tf.constant(list(range(4, -1, -1)))
selected = tf.gather(params, indices)

selected is now

[[24 25 26 27 28 29]
 [18 19 20 21 22 23]
 [12 13 14 15 16 17]
 [ 6  7  8  9 10 11]
 [ 0  1  2  3  4  5]]

---

If you need access to any other than the first dimension, you could work around that using tf.transpose: E.g. to gather columns instead of rows in our example, you could do this:

indices = tf.constant([0, 2])
selected = tf.gather(tf.transpose(params, [1, 0]), indices)
selected_t = tf.transpose(selected, [1, 0]) 

selected_t is of shape [5, 2] and reads:

[[ 0  2]
 [ 6  8]
 [12 14]
 [18 20]
 [24 26]]

However, tf.transpose is rather expensive, so it might be better to use tf.gather_nd for this use case.