Fix slice from view tensors (#129)

* Fix slice from view tensors

* Update test name

test/test_mps.py

@@ -1319,6 +1319,31 @@ def test_expand_cpu_to_mps_copy(self):
 
  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')