如何計算 Tensorflow 中矩陣對角線的總和？

Question

說，我有矩陣4x4像：

1 2 3 4
5 6 7 8
4 3 2 1
8 7 6 5`

我想用以下元素獲得矩陣2*4-1 ：

8
4+7
5+3+6
1+6+2+5
2+7+1
3+8
4

我怎樣才能在 Tensorflow 中做到這一點？ 當然，使用張量 - 我有形狀為[a,b,c,..,l,n,n]張量並且想要獲得形狀為[a,b,c,...,l,2*n-1] [a,b,c,..,l,n,n]張量[a,b,c,...,l,2*n-1] 。 是否有任何單一的功能可以做到這一點或循環是唯一的方法？

Answer 1

您可以使用tf.py_func包裝一個numpy函數。

import tensorflow as tf
import numpy as np

def np_all_trace_sum(a):
    n = a.shape[-1]
    all_trace_sum = [a.trace(i,axis1=-1,axis2=-2) for i in range(n-1,-n,-1)] # shape = (2*n-1,a,b,c,..,l)
    return np.moveaxis(all_trace_sum,0,-1) # shape = (a,b,c,..,l,2*n-1)

A = tf.placeholder(shape=[None,None,4,4],dtype=tf.float32)
result = tf.py_func(np_all_trace_sum, [A], tf.float32)

a = np.array([[1,2,3,4],[5,6,7,8],[4,3,2,1],[8,7,6,5]])

with tf.Session() as sess:
    print(sess.run(result,feed_dict={A:[[a,a,a],[a,a,a]]}))

[[[ 8. 11. 14. 14. 10. 11.  4.]
  [ 8. 11. 14. 14. 10. 11.  4.]
  [ 8. 11. 14. 14. 10. 11.  4.]]

 [[ 8. 11. 14. 14. 10. 11.  4.]
  [ 8. 11. 14. 14. 10. 11.  4.]
  [ 8. 11. 14. 14. 10. 11.  4.]]]