python之如何计算 Tensorflow 中矩阵对角线的总和

比如说，我有矩阵 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]。是否有任何单一函数可以执行此操作或循环是唯一的方法？

请您参考如下方法:

您可以使用 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.]]]