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