Create indices for a 2D Tensor

Viewed 1829

I'm trying to do something like this:

Suppose the input Tensor is a (2, 3) Tensor with the value like:

[[1 2]
 [3 4]
 [5 6]]

I am flattening the 2D Tensor into 1D by doing:

input = tf.reshape(input, [-1])

So the input now becomes [1 2 3 4 5 6].

But I also create a 1D Tensor that indexes the previous Tensor, so the desired output is [0 0 1 1 2 2]. How should I create this Tensor in TF?

In general, if the input Tensor has a shape of (X, Y). I would like to create a 1D Tensor that looks like:

[0 0 0 ... 0 1 ....1 2 ... 2 ... X-1 ... X-1]

where each value is repeated Y - 1 times.

1 Answers
Related