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