numpy数组布尔索引修改的替代方案、Torch M1适配及效率分析

MPS设备布尔索引赋值报错的临时修复 #

你遇到的operator 'aten::_index_put_impl_' is not implemented for MPS是PyTorch MPS后端当前的API覆盖限制，以下两种方法可以绕过：

方法1：CPU中转赋值 #

将张量临时转移到CPU完成赋值操作，再转回MPS：

mask = state == s
# 转移到CPU执行赋值
waiting_times_cpu = waiting_times.cpu()
new_vals_cpu = find_waiting_times(s, state[mask].shape).cpu()
waiting_times_cpu[mask.cpu()] = new_vals_cpu
# 同步回MPS设备
waiting_times.data.copy_(waiting_times_cpu.to(waiting_times.device))

注意：大张量跨设备拷贝会有性能损耗，适合小规模数据场景。

方法2：使用scatter_替代布尔索引赋值 #

scatter_是MPS支持的原地更新操作，先通过torch.where提取目标位置索引，再执行更新：

mask = state == s
indices = torch.where(mask)[0]  # 获取所有符合条件的位置索引
new_vals = find_waiting_times(s, indices.shape)
# 一维张量用dim=0，高维需对应调整dim参数
waiting_times.scatter_(dim=0, index=indices, src=new_vals)

这个方法无跨设备拷贝，性能更稳定，是跨CUDA/MPS/CPU的通用方案。

数组子集编辑的通用模式与效率权衡 #

1. 布尔索引赋值（tensor[mask] = vals） #

• 优势：代码简洁易读，CUDA/CPU上优化成熟，内存访问连续时速度极快；
• 劣势：MPS当前不支持，大规模张量的布尔掩码会占用额外内存；
• 适用场景：CUDA/CPU环境下，逻辑简单的子集更新。

2. scatter_/scatter操作 #

• 优势：全设备兼容（MPS/CUDA/CPU），原地操作内存开销小，索引定位精准；
• 劣势：需要先提取目标索引，高维张量需指定正确的dim参数；
• 适用场景：跨设备兼容需求，或MPS环境下的子集更新。

3. torch.where条件替换 #

• 优势：全设备支持，无分支向量化操作，适合需要保留原张量非目标位置值的场景；
• 劣势：需将更新值扩展为原张量同形状，额外占用内存；
• 示例：

  mask = state == s
  new_vals = find_waiting_times(s, mask.sum().shape)
  # 扩展更新值到原张量形状
  expanded_vals = torch.zeros_like(waiting_times)
  expanded_vals.scatter_(0, torch.where(mask)[0], new_vals)
  # 生成新张量（原地操作可用waiting_times.where_(~mask, expanded_vals)）
  waiting_times = torch.where(mask, expanded_vals, waiting_times)
  

4. 原地vs非原地操作对比 #

• 原地操作（如scatter_、[mask] =）：内存开销低，无需额外存储新张量，但依赖设备对操作的支持；
• 非原地操作（如torch.where）：兼容性更强，但会生成新张量，内存开销大。

长期建议 #

关注PyTorch官方MPS后端的更新，_index_put_impl_的实现缺口会在后续版本中修复。当前优先采用scatter_作为跨设备兼容的替代方案。

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