Получить строки тензора, удовлетворяющие условию тензорного потока

Question

input_mb = tf.placeholder(tf.int32, [None, 166, 1], name="input_minibatch")

Допустим, есть приведенный выше код. Я хочу получить строки указанного выше тензорного минибата так, чтобы первый элемент каждой найденной строки == a. Как мне это сделать в Tensorflow? Кроме того, как ты делаешь это в Numpy?

Answer 1

(Дано значение)

Чтобы достичь этого в numpy , вам просто нужно написать:

selected_rows = myarray[myarray[:,0]== a]

В тензор потока , используйте tf.where:

mytensor[tf.squeeze(tf.where(tf.equal(mytensor[:,0],a), None, None))