如何对Numpy数组/PyTorch张量不同行数做向量化求和、取最值？

首先给出原数组定义：

import numpy as np

M = np.array([
    [[1, 10, 100],
     [2, 20, 200],
     [3, 30, 300],
     [4, 40, 400]],        
    [[5, 50, 500],
     [6, 60, 600],
     [7, 70, 700],
     [8, 80, 800]]
])

一、向量化实现指定行求和 #

核心思路是构造布尔掩码标记需要参与求和的行，通过广播机制将掩码应用到数组的第三维度，最后沿行维度（axis=1）求和并保留维度。

1. 构造匹配需求的布尔掩码：

# 形状为(2,4)，对应M的前两个维度
mask = np.array([
    [False, True, True, True],  # 第一个子数组：第2、3、4行（索引1-3）参与求和
    [True, True, False, False]   # 第二个子数组：第1、2行（索引0-1）参与求和
])

2. 应用掩码并求和：

# 将掩码扩展为(2,4,3)以匹配M的形状，逐元素相乘后沿行维度求和
sum_result = (M * mask[:, :, np.newaxis]).sum(axis=1, keepdims=True)
print(sum_result)

输出结果与需求一致：

[[[  9  90 900]]

 [[ 11 110 1100]]]

二、向量化实现指定行求最小值/最大值 #

求和的掩码乘法逻辑不适用于极值计算（会引入无效值干扰结果），这里改用np.where将不需要的行替换为不影响极值的哨兵值，再沿行维度计算极值。

求最小值 #

# 把不需要的行替换为无穷大（inf），取min时只会保留有效行的值
min_masked = np.where(mask[:, :, np.newaxis], M, np.inf)
min_result = min_masked.min(axis=1, keepdims=True)
print(min_result)

输出：

[[[  2  20 200]]

 [[  5  50 500]]]

求最大值 #

# 把不需要的行替换为负无穷（-inf），取max时只会保留有效行的值
max_masked = np.where(mask[:, :, np.newaxis], M, -np.inf)
max_result = max_masked.max(axis=1, keepdims=True)
print(max_result)

输出：

[[[  4  40 400]]

 [[  6  60 600]]]

内容的提问来源于stack exchange，提问作者Sengiley