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.]]]