PyTorch MPS后端中where操作的非连续条件张量问题分析

在PyTorch的MPS后端实现中，当使用where操作时，如果条件张量(cond)是非连续的(contiguous)，可能会导致计算结果不正确。这个问题在最新版本的PyTorch nightly构建中被发现。

问题现象

当在MPS设备上执行where操作时，如果条件张量经过转置操作变为非连续状态，计算结果会与CPU上的结果不一致。具体表现为：

1. 在CPU上，where操作能够正确根据条件张量选择目标张量或零值
2. 在MPS上，当条件张量非连续时，where操作会错误地全部返回零值

技术背景

where操作是PyTorch中的一个重要函数，其原型为torch.where(condition, x, y)，它会根据condition的真假从x或y中选择元素。在底层实现中，这个操作需要高效处理各种张量布局，包括连续和非连续的内存布局。

MPS(Apple Metal Performance Shaders)是PyTorch为苹果设备提供的后端，专门优化了在苹果芯片上的张量运算性能。然而，在某些边界情况下，其实现可能与CPU后端存在差异。

问题复现

通过以下代码可以稳定复现该问题：

device = "mps"
diff = torch.tensor([[True, True], [True, True]], dtype=torch.bool)
diff = diff.T  # 使张量变为非连续
target = torch.tensor([[0, 0], [0, 1]])

# CPU结果正确
rcpu = torch.where(diff, target, 0)  # 输出[[0,0],[0,1]]

# MPS结果错误
diffmps = diff.to(device)
targetmps = target.to(device)
rmps = torch.where(diffmps, targetmps, 0)  # 输出[[0,0],[0,0]]

影响范围

这个问题主要影响：

1. 使用MPS后端的苹果设备用户
2. 涉及非连续条件张量的where操作场景
3. 可能影响依赖where操作的更高级操作，如某些分解操作(decompositions)

解决方案

对于遇到此问题的用户，可以采取以下临时解决方案：

1. 在执行where操作前，确保条件张量是连续的，可以使用.contiguous()方法
2. 暂时使用CPU进行计算，待问题修复后再切换回MPS

PyTorch团队已经确认此问题并标记为需要修复，预计在未来的版本更新中会解决这个非连续张量处理的边界情况。