Tensorflowで行列積の対角を計算する方法は？

Question

AとBの形状が(M, N)の非常に大きいMと小さいNの2つの行列を持っています。私はそれらを乗算し、その結果の対角線利用したいと思います

は：

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

残念ながら、これはメモリに収まらないことができる非常に大きな(M, M)マトリックスの作成を必要とします。

しかし、このデータのほとんどは必要ありません。ですから、この値を1ステップで計算することは可能ですか？

einsumがありますが、合計がありませんか？

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

か：：

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

---

例：あなたが必要なもの

Answer 1

は同等です

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

あなたはを使用することができますこれを得るためAとB transposeの：

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

コード：

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]