Hi  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 4. 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 !