Add additional tests for view slice tensors (#86282)
Fixes https://github.com/pytorch/pytorch/issues/83995 and https://github.com/pytorch/pytorch/issues/84489
Pull Request resolved: https://github.com/pytorch/pytorch/pull/86282
Approved by: https://github.com/kulinseth
diff --git a/test/test_mps.py b/test/test_mps.py
index 77c2a34..75b3c35 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -1319,6 +1319,31 @@
self.assertEqual(x_cpu, x.cpu())
+ def test_view_slice(self):
+ # https://github.com/pytorch/pytorch/issues/83995
+ NUM_SAMPLES = 60
+ s = (0, 1)
+
+ X = torch.rand(8000, 3, dtype=torch.float32, device='cpu')
+ X_mps = X.detach().clone().to("cpu")
+
+ idx = torch.randint(0, X.shape[0], (1,)).repeat(len(s))
+ pts = torch.randint(0, X.shape[0], (NUM_SAMPLES, X.shape[1]))
+ idx_mps = idx.to("mps")
+ pts_mps = pts.to("mps")
+ pts[:, s] = idx
+ pts_mps[:, s] = idx_mps
+
+ actual_pts = torch.zeros(NUM_SAMPLES, X.shape[1], dtype=torch.float)
+ actual_pts_mps = torch.zeros(NUM_SAMPLES, X.shape[1], dtype=torch.float, device="mps")
+
+ for i in range(NUM_SAMPLES):
+ for j in range(X.shape[1]):
+ actual_pts_mps[i, j] = X_mps[pts_mps[i, j], j]
+ actual_pts[i, j] = X[pts[i, j], j]
+ self.assertEqual(actual_pts[i, j], actual_pts_mps[i, j])
+
+
def test_slice(self):
values = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]
cpu_x = torch.tensor(values, device='cpu')
