张量流中张量点积的有效实现

Question

我有三个x,P,A大小(K,N), (K,M), (M,N,N)的张量。我想计算一个大小为Y的矩阵(K,N)，以便：

目前我在tensorflow中使用以下实现：

PA=tf.einsum('ij,jkl->ikl',P,A)
Y= tf.einsum('ij,ijk->ik', x,PA)

我想知道是否有任何其他有效的程序来实现这一点。我尝试使用内置矩阵操作，例如matmul()和batch_matmul()，但我无法找到使用它们的解决方案。任何有用的建议都会有所帮助。谢谢！