How to slice tensorflow tensor differently for each row at once?

user_4474068 0
Viewed 305

I have simple tensor

a = tf.constant([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])

and want to slice it, but I need to do it differently for each of the rows. This slice operation is described by another tensor

b = tf.constant([[0, 1], [2, 4], [2, 5]])

It means that from the first row of tensor a I need elements from 0 to 1, from the second row from 2 to 4 and so on. So the final final result will be

[
  [1],
  [8, 9],
  [13, 14, 15]
]

My first idea was to fill ranges between begin and end of a slice, but unfortunately, doing it with map_fn is not possible because result's rows have different lengths.

Does anyone know how to do such operation?

data-science tensorflow python
2 Answers
  • Basically we have two arrays to iterate. One with actual data, other with range to return.
  • Therefore, zip function can help iterate over elements from multiple arrays one by one.
import tensorflow as tf

a = tf.constant([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])
b = tf.constant([[0, 1], [2, 4], [2, 5]])

# As you iterate, provided a and b have same length
# [1, 2, 3, 4, 5] sliced as [0:1]
# [6, 7, 8, 9, 10] sliced as [2:4]
# [11, 12, 13, 14, 15] sliced as [2:5]

[data.numpy().tolist()[start:end] for data, (start, end) in zip(a,b)]

Output:

[[1], [8, 9], [13, 14, 15]]
user_28650570

If the size of b is known at graph compile time, then you can slice each row separately.

import tensorflow as tf

a = tf.constant([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])
b = tf.constant([[0, 1], [2, 4], [2, 5]])

r = []
for i in range(3):
    bi = b[i]
    r.append(a[i][bi[0]: bi[1]])
print(r)
user_34152350
Related