Блог

Например, у меня есть тензор формы [2, 2] . Я хочу получить 3D-тензор индексов: [[[0, 0], [0, 1]], [[1, 0], [1, 1]]].

Я использую следующий подход:

 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)
 

Есть ли лучший подход?

Я бы поступил так же, как в numpy: meshgrid вероятно, это правильный путь (в сочетании с stack получением 1 массива).

 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)>