[MPS] Raise NotImplementedError for pooling ops that require 64-bit indexing (#194460)
Continuation of #194373
More errors. Below script causes failures:
```
import torch
import torch.nn.functional as F
def run(name, shape, op, bwd):
x = torch.randn(shape, dtype=torch.half)
xm = x.to("mps").requires_grad_(bwd)
xc = x.float().requires_grad_(bwd)
try:
ym, yc = op(xm), op(xc)
if bwd:
g = torch.randn_like(yc)
ym.backward(g.to("mps").half())
yc.backward(g)
rm, rc = xm.grad.cpu().float(), xc.grad
else:
rm, rc = ym.detach().cpu().float(), yc.detach()
bad = int(((rm - rc).abs() > 0.1).sum())
print(f"{name}: {'FAIL' if bad else 'ok'} mismatched={bad}")
except NotImplementedError as e:
print(f"{name}: guarded ({e})")
torch.mps.empty_cache()
P2 = (34, 1, 8192, 8192)
P3 = (70, 32, 100, 100, 100)
run("max/adaptive fwd template", P2, lambda t: F.max_pool2d(t, 2), False)
run("max bwd template", P2, lambda t: F.adaptive_max_pool2d(t, 128), True)
run("avg fwd template", P2, lambda t: F.avg_pool2d(t, 2, ceil_mode=True), False)
run("avg bwd template", P3, lambda t: F.avg_pool3d(t, 2), True)
```
leads to:
```
max/adaptive fwd template: FAIL mismatched=32208577
max bwd template: FAIL mismatched=30132
avg fwd template: FAIL mismatched=28232769
avg bwd template: FAIL mismatched=39205174
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/194460
Approved by: https://github.com/malfet I
Isalia20 committed
84c32f8d959c5226f3b85d12d4a8ba260e13b52a
Parent: 20ef2cc
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
on 8/24/2026, 1:20:47 PM