CODEWITHSUNDEEP
pythontensorflow
LOST
LOST

Reputation: 3250

How to get multiple slices of the same size from a `Tensor`?

I have a tensor I of indices (of unknown shape n), that indicates where slices each of constant length slice_length should start from a tensor T (of shape t_length).

What I want to get is a tensor of shape (n, slice_length), consisting of slices of T.

For example, if T is [0,1,2,3,4,5,6], I is [1,3,1], and slice_length is 2, the resulting tensor should be

[[1,2],
 [3,4],
 [1,2]]

What is the most efficient way to do it?

Upvotes: 1

Views: 123

Answers (2)

Innat
Innat

Reputation: 17219

Based on your comment, I understand that you need to first build your indices based on what you need, and the rest of the stuff is pretty simple. Here is one way to achieve what you mention in the comment box.

import tensorflow as tf 

T = tf.constant([0,1,2,3,4,5,6])
x = tf.constant([1,3,1])
y = x + 1
I = tf.stack([x, y], -1)

tf.gather(T, I)
<tf.Tensor: shape=(3, 2), dtype=int32, numpy=
array([[1, 2],
       [3, 4],
       [1, 2]], dtype=int32)>

Upvotes: 0

LOST
LOST

Reputation: 3250

Here's what I ended up doing. Not sure if there is a better way.

slice_base_indices = tf.range(0, slice_length)
slice_base_indices = tf.expand_dims(slice_base_indices, axis=1)
indices = slice_base_indices + I
indices = tf.transpose(indices)
return tf.gather(T, indices)

Upvotes: 1

Related Questions