fix slice w/ dynamic shapes (#153131)

Summary: guard_size_oblivious has side effects that'll result in invalid strides when slice nodes take negative index on dynamic input shapes.
Cause overflow error with a huge number “9223372036854776048”
Test Plan: CIs should pass.

Differential Revision: D74354663

Pull Request resolved: https://github.com/pytorch/pytorch/pull/153131
Approved by: https://github.com/laithsakka
This commit is contained in:
Chao Gu
2025-06-13 15:53:17 +00:00
committed by PyTorch MergeBot
parent a5938ff431
commit 338a8c7853

View File

@ -750,11 +750,11 @@ def slice_forward(
elif guard_size_oblivious(start_val > sizes[dim]):
start_val = sizes[dim]
if guard_size_oblivious(end_val < start_val):
if statically_known_true(end_val == sys.maxsize):
end_val = sizes[dim]
elif guard_size_oblivious(end_val < start_val):
end_val = start_val
elif statically_known_true(end_val == sys.maxsize) or guard_size_oblivious(
end_val > sizes[dim]
):
elif guard_size_oblivious(end_val > sizes[dim]):
end_val = sizes[dim]
storage_offset = self.storage_offset() + start_val * strides[dim]