[MPS] Get the correct size of the view tensor when copying from cpu to mps (#81730)
Fixes: https://github.com/pytorch/pytorch/issues/81567, https://github.com/pytorch/pytorch/issues/80844
* Get the correct size of the view tensor when copying from cpu to mps
* Use 'computeStorageNbytesContiguous' to get the size just when src is a view
* Add asserts and tests to check for storage_offset
* Add testcase for https://github.com/pytorch/pytorch/issues/80844
* Replace assert_allclose with assertEqual
* Replace TORCH_CHECK with TORCH_INTERNAL_ASSERT
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81730
Approved by: https://github.com/razarmehr, https://github.com/albanD
diff --git a/test/test_mps.py b/test/test_mps.py
index c6e1063..12b5643 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -1490,6 +1490,58 @@
z = y.as_strided(size=(32, 3), stride=(1, 0)).to("cpu")
self.assertEqual(x.to("cpu").as_strided(size=(32, 3), stride=(1, 0)), z)
+ def test_type_casting(self):
+ # https://github.com/pytorch/pytorch/issues/81567
+ def helper(data, to_dtype):
+ a_cpu = torch.tensor(data)
+ a_mps = a_cpu.to(torch.device('mps'))
+
+ res_cpu = a_cpu.type(to_dtype)
+ res_mps = a_mps.type(to_dtype)
+ self.assertEqual(res_cpu, res_mps)
+
+ helper([9.0, 3.0, 5.0, 4.0], torch.LongTensor)
+ helper([9.0, 3.0, 5.0, 4.0], torch.FloatTensor)
+ helper([9.0, 3.0, 5.0, 4.0], torch.IntTensor)
+ helper([9.0, 3.0, 5.0, 4.0], torch.ShortTensor)
+ helper([9.0, 3.0, 5.0, 4.0], torch.HalfTensor)
+ helper([9.0, 3.0, 5.0, 4.0], torch.CharTensor)
+ helper([9.0, 3.0, 5.0, 4.0], torch.ByteTensor)
+
+ def test_to_casting(self):
+ # https://github.com/pytorch/pytorch/issues/81567
+ def helper(data, to_dtype):
+ a_cpu = torch.tensor(data)
+ a_mps = a_cpu.to(torch.device('mps'))
+
+ res_cpu = a_cpu.to(to_dtype)
+ res_mps = a_mps.to(to_dtype)
+ self.assertEqual(res_cpu, res_mps)
+
+ helper([9.0, 3.0, 5.0, 4.0], torch.int64)
+ helper([9.0, 3.0, 5.0, 4.0], torch.float)
+ helper([9.0, 3.0, 5.0, 4.0], torch.int32)
+ helper([9.0, 3.0, 5.0, 4.0], torch.short)
+ helper([9.0, 3.0, 5.0, 4.0], torch.half)
+ helper([9.0, 3.0, 5.0, 4.0], torch.int8)
+ helper([9.0, 3.0, 5.0, 4.0], torch.uint8)
+
+ def test_storage_offset_greater_than_src_nbytes(self):
+ # https://github.com/pytorch/pytorch/issues/80844
+ n_tensors = 100
+ n_tensor_elems = 784
+ elems = torch.arange(n_tensors * n_tensor_elems, dtype=torch.float32)
+
+ tensor_list = []
+ for i in range(0, n_tensors - 1):
+ # create a list of contiguous view tensors (view tensor created by the slice op)
+ t = elems[n_tensor_elems * i : n_tensor_elems * (i + 1)]
+ tensor_list.append(t)
+
+ for i in range(0, n_tensors - 1):
+ t = tensor_list[i].view(1, 784)
+ t_mps = t.to("mps")
+ self.assertEqual(t, t_mps.cpu())
class TestLogical(TestCase):
def _wrap_tensor(self, x, device="cpu", dtype=None, requires_grad=False):
