多维矩阵对角元素的批处理函数

Question

我有两个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中创建的对角矩阵，而不是以迭代的方式。实现这一点的等效操作集是什么？你知道吗