How to calculate diagonal of a matrix product in 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]