Блог

Учитывая такой целочисленный тензор x = [2 0 1 0] и константу C = 3 , я хотел бы построить тензор

 [[1 1 0],
 [0 0 0],
 [1 0 0],
 [0 0 0]]
 

Так C же как и количество столбцов, len(x) это количество строк. Каждая запись x указывает число 1, с которого должна начинаться строка, в то время как остальная часть этой строки должна быть заполнена 0.

Лучший код TensorFlow, который я придумал, это:

 x = tf.constant([2, 0, 1, 0])
C = 3
r = tf.map_fn(fn=lambda i: tf.concat([tf.ones(i, dtype=tf.int32),
                                      tf.zeros(C-i, dtype=tf.int32)], 
                                     axis=0), 
              elems=x)
 

Это работает нормально:

 In [6]: print(r)
tf.Tensor(
[[1 1 0]
 [0 0 0]
 [1 0 0]
 [0 0 0]], shape=(4, 3), dtype=int32)
 

Но есть ли простой способ достичь того же результата без tf.map_fn() каких-либо циклов или без них, но с векторизованными операциями, подобными tf.scatter_nd() или похожими?

tf.sequence_mask() дает желаемый результат:

 In [8]: s = tf.sequence_mask(x, maxlen=C, dtype=tf.int32)

In [9]: print(s)
tf.Tensor(
[[1 1 0]
 [0 0 0]
 [1 0 0]
 [0 0 0]], shape=(4, 3), dtype=int32)