我有两个像这样的numpy矩阵：

matrix_a = np.random.random((2, 3))
matrix_b = np.random.random((2, 3, 5))

使用以下函数，我为 matrix_a 的轴1中的每个元素生成一个对角矩阵：

def foo(M):
  b = np.zeros((M.shape[0], M.shape[1], M.shape[1]))
  diag = np.arange(M.shape[1])
  b[:, diag, diag] = M
  return b

我也在扩大 matrix_a 的尺寸，如下所示：

diag_mixed_alignments = np.expand_dims(diag_mixed_alignments, axis=1)

matrix_a 的形状将变为 (2, 1, 3, 3) .

假设第0维的大小2表示批量大小 . 我想为每个批次执行以下操作：

for row in range(0, diag_mixed_alignments.shape[0]):
      cur_M = diag_mixed_alignments[row, :]
      cur_V = values[row, :]

      intmd_op1 = np.sum(np.matmul(cur_V.T, cur_M), axis=2)

      op1.append(intmd_op1)

以上将产生一个形状矩阵 (2, 1, 5)

我的目标是在tensorflow中执行相同的操作，这样我就可以考虑在 matrix_a not in an iterative manner 中创建的对角矩阵 . 实现这一目标的等效操作集是什么？

python tensorflow matrix-multiplication