如何在Tensorflow中计算矩阵乘积的对角线？

Question

I have two matrices A and B of shape (M, N) with very large M and small N .

I would like to multiply them and then take diagonal of a result:

C = tf.matmul(A, B)
D = tf.diag_part(C)

Unfortunately, this requires of creating of very big (M, M) matrix, which can't fit into memory.

But most of this data I don't need. So, is it possible to calculate this value in one step?

Is there something like einsum but without summing?

Answer 1

What you need is equivalent to:

tf.einsum('ij,ij->i', A, B)

or:

tf.reduce_sum(A * B, axis=1)

---

Example :

A = tf.constant([[1,2],[2,3],[3,4]])
B = tf.constant([[3,4],[1,2],[2,3]])

with tf.Session() as sess:
    print(sess.run(tf.diag_part(tf.matmul(A, B, transpose_b=True)))) 
# [11  8 18]

with tf.Session() as sess:
    print(sess.run(tf.reduce_sum(A * B, axis=1)))
#[11  8 18]

with tf.Session() as sess:
    print(sess.run(tf.einsum('ij,ij->i', A, B)))
#[11  8 18]

Answer 2

You can use the dot product of A and B transpose to obtain the same:

tf.reduce_sum(tf.multiply(A, tf.transpose(B)), axis=1)

The code:

import tensorflow as tf
import numpy as np

A = tf.constant([[1,4, 3], [4, 2, 6]])
B = tf.constant([[5,4,],[8,5], [7, 3]])

E = tf.reduce_sum(tf.multiply(A, tf.transpose(B)), axis=1)

C = tf.matmul(A, B)
D = tf.diag_part(C)
sess = tf.InteractiveSession()

print(sess.run(D))
print(sess.run(E))

#Output
#[58 44]
#[58 44]