我正在尝试使用作为变量进行numpy.matmul调用

• 尺寸为(p, t, q)的矩阵A
• 尺寸为(r, t)的矩阵B。
• 形状r和p类别的categories向量，用于获取B的切片并定义A的索引。

使用每个类别的索引迭代地进行乘法。对于每个类别p_i，我从A中提取子矩阵（t, q）。然后，我将它们与B的子集相乘（x, t），其中x是由r == p_i定义的掩码。(x, t)和(t, q)的矩阵乘法产生存储在S[x]处的输出(x, q)。
我已经注意到我无法计算出这个算法的非迭代版本。第一个片段描述了一个迭代解决方案。第二个片段是我希望得到的尝试，其中所有内容都作为一个单步计算，可能会更快。然而，它是不正确的，因为矩阵A是三维而不是二维的。也许没有办法在NumPy中通过一个调用来完成这一点。总的来说，寻找建议/想法来尝试。
谢谢!

import numpy as np
p, q, r, t = 2, 9, 512, 4
# data initialization (random)
np.random.seed(500)
S = np.random.rand(r, q)
A = np.random.randint(0, 3, size=(p, t, q))
B = np.random.rand(r, t)
categories = np.random.randint(0, p, r)

print('iterative') # iterative
for i in range(p):
    # print(i)
    a = A[i, :, :]
    mask = categories == i
    b = B[mask]
    print(b.shape, a.shape, S[mask].shape,
          np.matmul(b, a).shape)
    S[mask] = np.matmul(b, a)

print(S.shape)

一个简单的方法来写下来

S = np.random.rand(r, q)
print(A[:p,:,:].shape)
result = np.matmul(B, A[:p,:,:])

# iterative assignment
i = 0
S[categories == i] = result[i, categories == i, :]
i = 1
S[categories == i] = result[i, categories == i, :]

下一个代码段将在乘法步骤中产生错误。
一个二个一个一个