0

TL;DR: TensorFlow tensor is of shape (50, 50, 6), want these indices (:, :, (0, 2, 3)). How to extract them?

Here is an example array I am working with:

import numpy as np

a = np.random.randint(0,10, (50, 50, 6))

I want to extract the the first, third, and fourth row; in other words I need all these entries (:, :, (1, 3)), which works for numpy arrays:

out = a[:,:, [0, 2, 3]]
out.shape #(50, 50, 3)

Working with a tensor t = tf.convert_to_tensor(a) and then calling the index like

t[:,:, [0, 2, 3]]

throws an error:

TypeError: Only integers, slices (`:`), ellipsis (`...`), tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid indices, got [0, 1, 3]

For numpy I have found the following relevant questions, but they naturally focus on numpy arrays:

How to slice a 2D array non-consecutively in Python

Slicing a numpy array along a dynamically specified axis

Looking at the TF documentation I found gather_nd and boolean_mask, which I feel are helpful, but I must freely admit that I have not understood the docs at this part. On SO I found this question How to select elements of a tensor along a specific axis in TensorFlow, which focuses on single elements; I am looking for complete dimensions (if that's the right wording here).

How can I do the numpy thing in TensorFlow?

emil
  • 194
  • 1
  • 11
  • Can you extract every channel by itself and combine them afterwars? Or convert it to a numpy array and then use `out = a[:,:, [0, 2, 3]]` as before? – chillking Apr 21 '21 at 14:14
  • If you want the 0,2 and 3rd element of the last axis in the tensor, you can use tf.gather as follows: tf.gather(t,indices=[0, 2, 3],axis=-1)) – Abhilash Rajan Apr 21 '21 at 14:38
  • @AbhilashRajan your solution is what I've been looking for! If you don't mind then you can write your comment as a complete answer and I'll gladly accept it. – emil May 06 '21 at 11:22

1 Answers1

1

If you want the 0,2 and 3rd element of the last axis in the tensor, you can use tf.gather as follows: tf.gather(t,indices=[0, 2, 3],axis=-1))

Abhilash Rajan
  • 349
  • 1
  • 7