CODEWITHSUNDEEP
pythontensorflow
Andrey
Andrey

Reputation: 6377

How to get a tensor of indices in tensorflow?

For example, I have a tensor of shape [2, 2]. I want to get a 3D tensor of indices: [[[0, 0], [0, 1]], [[1, 0], [1, 1]]].

I am using the following approach:

my_shape = [2, 2]
my_range = my_shape[0]*my_shape[1]
m = tf.range(0, my_range)
def get_inds(t, last_dim):
  return tf.convert_to_tensor([t // last_dim, t % last_dim])

inds = tf.map_fn(fn=lambda t: get_inds(t, my_shape[-1]), elems=m)
sh = tf.concat([my_shape, [2]], -1)
inds = tf.reshape(inds, sh)

Is there any better approach ?

Upvotes: 1

Views: 292

Answers (1)

Lescurel
Lescurel

Reputation: 11651

I would do like in numpy : meshgrid is probably the way to go (combined with stack to get 1 array).

import tensorflow as tf
shape = (2,2)
x = tf.range(shape[0])
y = tf.range(shape[1])
xx,yy = tf.meshgrid(x,y)
indices = tf.stack([yy,xx],axis=2)

>>> indices
<tf.Tensor: shape=(2, 2, 2), dtype=int32, numpy=
array([[[0, 0],
        [0, 1]],

       [[1, 0],
        [1, 1]]], dtype=int32)>

Upvotes: 1

Related Questions