tensorflow:将矩阵的某些行与另一个矩阵中的某些列相乘

tf.gather 具体如下:

import tensorflow as tf
a=tf.constant([[1,2,3],[4,5,6],[7,8,9]])
b=tf.constant([[1,0,1],[1,0,2],[3,3,-1]])

#taking rows 0,1 from a, and columns 0,2 from b
ind_a=tf.constant([0,1])
ind_b=tf.constant([0,2])

r_a=tf.gather(a,ind_a)

#tf.gather access the rows, so we use it together with tf.transpose to access the columns
r_b=tf.transpose(tf.gather(tf.transpose(b),ind_b))

# the diagonal elements of the multiplication
res=tf.diag_part(tf.matmul(r_a,r_b))
sess=tf.InteractiveSession()
print(r_a.eval())
print(r_b.eval())
print(res.eval())

这张照片

#r_a
[[1 2 3]
 [4 5 6]]

#r_b
[[ 1  1]
 [ 1  2]
 [ 3 -1]]

#result
[12  8]