Fix overflow in checkInBoundsForStorage (#147352)

Use `computeStorageNbytes` (which checks for overflows) to include the computation re the storage_offset

Pull Request resolved: https://github.com/pytorch/pytorch/pull/147352
Approved by: https://github.com/albanD
This commit is contained in:
Mikayla Gawarecki
2025-02-24 17:12:53 -08:00
committed by PyTorch MergeBot
parent 6ccbff1450
commit e64441915f
2 changed files with 16 additions and 3 deletions

View File

@ -1995,6 +1995,18 @@ class TestOldViewOps(TestCase):
with self.assertRaisesRegex(RuntimeError, "Stride calculation overflowed"):
x.resize_([0, 4, 2305843009213693952])
@onlyNativeDeviceTypes
def test_as_strided_overflow_storage_offset(self, device):
t = torch.randn(2, 3, device=device)
with self.assertRaisesRegex(
RuntimeError, "Storage size calculation overflowed"
):
torch.as_strided(t, [1], [1], 2**63 - 1)
with self.assertRaisesRegex(
RuntimeError, "Storage size calculation overflowed"
):
torch.as_strided(t, [1], [1], 2**61 - 1)
def test_view_all_dtypes_and_devices(self, device):
for dt in all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool):
x = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=dt, device=device)