Using log-sum-exp for matrix multiplication

General Discussion
#1

Hi :slight_smile: I am trying to use TensorFlow perform a matrix multiplication of log matrices using a log-sum-exp operation. Very shortly put: I have a tensor A of shape (i,k) and a tensor B of shape (k,j) and want to perform the operation

C[i,j] = log ( sum_k ( exp( A[i, k] + B[k, j]) ) ),

as in eq. 11 of https://arxiv.org/pdf/1904.04676.pdf. Simply adding the tensors A and B is of course not possible, as they have to be added in a very specific manner that is almost interwoven with the log-sum-exp operation itself. Something like tf.math.logsumexp(A+B) will not work because the shapes of the tensors do not match, for example.

My question therefore is: if this is not the correct way of coding the desired formula in TensorFlow (which it isn’t), what is?

Many thanks in advance :relaxed:!

#3

Hi @bstienen, When adding two input values of different shapes, Add follows NumPy broadcasting rules. The two input array shapes are compared element-wise. Starting with the trailing dimensions, the two dimensions either have to be equal or one of them needs to be 1.

For example,

x = np.ones(6).reshape(1, 2, 1, 3)
y = np.ones(6).reshape(2, 1, 3, 1)
a=tf.add(x, y).shape.as_list() #output:[2, 2, 3, 3]

Then you can perform tf.math.reduce_logsumexp(a). Thank You.