如何在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]的张量。有没有单个函数可以执行此操作，否则循环是唯一方法？

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