[MPS] Fix type casting copy with storage offset (#95573)
This PR handles the case where the `dst` tensor of type casting has a storage offset by creating a temporary buffer to store results and then copy them back to the dst with the offset added.
Fixes #95417
Pull Request resolved: https://github.com/pytorch/pytorch/pull/95573
Approved by: https://github.com/kulinseth
diff --git a/test/test_mps.py b/test/test_mps.py
index 96d36ef..a49b82b 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -2665,6 +2665,16 @@
y.permute(3, 2, 1, 0)[1::, ::2] = z
self.assertEqual(x, y.to('cpu'))
+ # See https://github.com/pytorch/pytorch/issues/95417
+ def test_copy_storage_offset(self):
+ x_cpu = torch.zeros(5, device="cpu", dtype=torch.float32)
+ x_mps = torch.zeros(5, device="mps", dtype=torch.float32)
+ update_cpu = torch.tensor([1, 1], device="cpu", dtype=torch.int64)
+ update_mps = torch.tensor([1, 1], device="mps", dtype=torch.int64)
+ x_cpu[2:4] = update_cpu
+ x_mps[2:4] = update_mps # implicit type casting and copy
+ self.assertEqual(x_cpu, x_mps)
+
# See https://github.com/pytorch/pytorch/pull/84742
# and https://github.com/pytorch/pytorch/pull/78319
def test_binops_dtype_precedence(self):
