Tensorflow聚集在具有不同形状的输入上

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)