从Tensorflow中的tf.matmul（tf.transpose（A），A）中获取对角线元素

Question

I have a matrix A with several columns, and I need to compute the self dot products, ie tf.matmul(A[:,i]), A[:,i], transpose_a=True) where i indexes columns of A . One way is to just compute tf.matmul(A,A, transpose_a=True) and extract the diagonal, but this involves a lot of superfluous multiplication (all non-diagonal results are discarded, and the non-diagonal results are most of the results for any matrix larger than 2 x 2). Another way is to do something like

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

and then collect out into a tf.Tensor . But it seems like this is a fairly common computation, so I would expect a function to exist for it (ie computing squared norms of weight vectors).

Answer 1

You can multiply the matrix by itself elementwise and then sum along axis 0.

import tensorflow as tf
tf.InteractiveSession()

A = tf.reshape(tf.range(12), (3, 4))
tf.reduce_sum(tf.pow(A, 2), axis=0).eval()

returns

array([ 80, 107, 140, 179], dtype=int32)

Answer 2

Just a little bit of linear algebra and you have your solution:

Which means that you need to do: tf.reduce_sum(tf.square(A), axis=0)