Make Tensor.set_ validate storage_offset when sizes/strides are unchanged (#147354)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/147354
Approved by: https://github.com/albanD
ghstack dependencies: #147352
This commit is contained in:
Mikayla Gawarecki
2025-02-26 19:37:28 -08:00
committed by PyTorch MergeBot
parent e64441915f
commit 536bce5a04
6 changed files with 77 additions and 13 deletions

View File

@ -7166,6 +7166,18 @@ class TestTorch(TestCase):
f_cpu = torch.randn((2, 3), dtype=torch.float32)
d_cpu = torch.randn((2, 3), dtype=torch.float64)
storage_offset = 0x41414141
with self.assertRaisesRegex(RuntimeError, "out of bounds for storage of size"):
t = torch.randn(1)
t.set_(t.untyped_storage(), storage_offset, t.size())
# if size changes, set_ will resize the storage inplace
t = torch.randn(1)
size = torch.Size([2, 3])
t.set_(t.untyped_storage(), storage_offset, size)
self.assertEqual(t.storage_offset(), storage_offset)
self.assertEqual(t.untyped_storage().nbytes(), (storage_offset + size[0] * size[1]) * 4)
# change dtype
self.assertRaises(RuntimeError, lambda: f_cpu.set_(d_cpu.storage()))
self.assertRaises(RuntimeError,