python - 从 Tensorflow 中的 tf.matmul(tf.transpose(A), A) 获取对角线元素

Question

我有一个包含几列的矩阵A，我需要计算自点积，即tf.matmul(A[:,i]), A[:,i ], transpose_a=True) 其中 i 对 A 的列进行索引。一种方法是只计算 tf.matmul(A,A, transpose_a=True) 并提取对角线，但这涉及大量多余的乘法(所有非对角线结果都被丢弃，并且非对角线结果被丢弃)。 - 对角线结果是任何大于 2 x 2 的矩阵的大部分结果。另一种方法是做类似的事情

out = []
for i in range(tf.shape(A)[1]):
    out.append(tf.matmul((A[:,i],A[:,i],transpose_a=True))

然后将out收集到tf.Tensor中。但这似乎是一个相当常见的计算，所以我希望它存在一个函数(即计算权重向量的平方范数)。

Answer 1

只需一点线性代数，你就可以得到解决方案:

这意味着您需要执行以下操作:tf.reduce_sum(tf.square(A), axis=0)