[MPS] Fix base shape size for view ops in case of multiple slices (#85934)
Fixes https://github.com/pytorch/pytorch/issues/84364, https://github.com/pytorch/pytorch/issues/85592
Fixes bug for view ops where the base shape would be incorectly determined.
E.g for the following tensor `torch.tensor([0.5, 0.5], device="mps")[1][None]`, we could consider the base shape of the parent tensor as 1, while the actual base shape is 2.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85934
Approved by: https://github.com/kulinseth
diff --git a/aten/src/ATen/native/mps/operations/View.mm b/aten/src/ATen/native/mps/operations/View.mm
index 9f8bc54..263bfd0 100644
--- a/aten/src/ATen/native/mps/operations/View.mm
+++ b/aten/src/ATen/native/mps/operations/View.mm
@@ -201,7 +201,9 @@
// IntArrayRef wouldn't own the data, so we use a static storage
static const int64_t shape_1d = 1;
// self.sizes().size() could be zero
- base_shape = self.sizes().size() ? self.sizes() : IntArrayRef(&shape_1d, 1);
+ base_shape = self.sizes().size() ? self.sizes() :
+ self.is_view() ? self._base().sizes() : IntArrayRef(&shape_1d, 1);
+
// base_shape will be retained in MPSAllocator until buffer gets recycled
if (self.storage().data())
set_buffer_shape(self.storage().data(), base_shape);
diff --git a/test/test_mps.py b/test/test_mps.py
index c251e92..fa5f663 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -1373,6 +1373,18 @@
for op in ["<=", "<", ">=", ">", "==", "!="]:
helper(op)
+ def test_slice_of_slice(self):
+ x = torch.tensor([0.5, 0.5], device="cpu")
+ x_mps = torch.tensor([0.5, 0.5], device="mps")
+
+ tensor = x[1][None]
+ tensor_mps = x_mps[1][None]
+
+ res = tensor.ne(0)
+ res_mps = tensor_mps.ne(0)
+
+ self.assertEqual(res, res_mps)
+
def test_index_storage_offset(self):
# https://github.com/pytorch/pytorch/issues/78107
