我有两个大小为(16,8400(的矩阵A，大小为(16/50400(的矩阵B；其中16是批量大小。
现在我需要做的是：
对于A中的每一批，我想计算一批A中每一行与一批B中每一行都的内积，然后得到最大内积值。换句话说，对于每一批，我在左边(a(有一个(8400(矩阵，在右边(B(有一(50400(矩阵。现在，对于A中的每一行(8*400(，我想逐个计算该行与B中的每行(1*400(的内积，并得到最大内积。因此，内积运算将产生(8*1(大小的向量，然后得到8个值中的最大值。对于每一批，它都会产生(16，50，1(大小的矩阵，这是我输出的首选维度。
数学公式如图所示在此处输入图像描述
我该怎么做？我不想使用for循环，因为那样效率会很低。出于同样的原因，我也不想使用重复方法。
这行得通吗？
C = np.max( np.matmul(A, B.transpose(0,2,1)) , axis=1 )