从 TensorFlow 中的 1 和 0 动态构建张量

Question

Given an integer tensor like x = [2 0 1 0] and a constant C = 3 I would like to build a tensor

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

So C is the number of columns, len(x) is the number of rows. Each entry of x specifies a number of 1s that a row should start with, while the remainder of that row is to be filled up with 0s.

The best TensorFlow code that I've come up with is this:

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)

This works OK:

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

But is there a straightforward way to achieve the same result without tf.map_fn() or any loops, but with vectorized ops like tf.scatter_nd() or similar?

Answer 1

tf.sequence_mask() produces the desired result:

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)