Tensorflow gather_nd 对具有不同形状的输入

Question

我需要根据另一个张量的值沿 axis=1 提取子张量。

state = tf.placeholder(shape=[None, None, 10], dtype=tf.float32)
length = tf.placeholder(shape=[None], dtype=tf.int32)

# this won't work, just be put here to demonstrate what I need
next_init_state = state[:, length - 1, :]

如果state和length具有确定性形状，那么next_init_state可以通过gather_nd

导出

state = tf.placeholder(shape=[10, 10, 10], dtype=tf.float32)
length = tf.placeholder(shape=[10], dtype=tf.int32)

index = tf.stack([tf.range(0, 10), length])
next_init_state = tf.gather_nd(state, index)

然而，由于状态和长度在我遇到的问题中都具有不确定的形状 None，gather_nd 方法将不起作用。至少我想不出让它发挥作用的方法。有什么办法可以解决吗？

Answer 1

我意识到这实际上是由Tensorflow的高阶函数解决的。

tf.map_fn(lambda x: x[0][x[1], :], (state, length), dtype=tf.float32)

Tensorflow gather_nd 对具有不同形状的输入

Tensorflow gather_nd over input with varied shape

python

tensorflow

tensor