From 1e0a59aca48f6855aec00e20c3339beee1197d3e Mon Sep 17 00:00:00 2001 From: Xingyu Date: Fri, 23 May 2025 01:54:13 +0800 Subject: [PATCH] fix: handle buffer size calculation in to_movement_ops and add scalar assignment test in torch_backend (#10464) --- extra/to_movement_ops.py | 2 +- extra/torch_backend/test.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/extra/to_movement_ops.py b/extra/to_movement_ops.py index 81afd306d8..68d4ef3dfe 100644 --- a/extra/to_movement_ops.py +++ b/extra/to_movement_ops.py @@ -41,7 +41,7 @@ def to_movement_ops(st: ShapeTracker) -> List[Tuple[MovementOps, Tuple]]: real_real_shape = [s for s,st in zip(real_shape, v.strides) if st] strides: List[int] = [abs(st) if isinstance(st,int) else st for st in v.strides if st] buffer_size = sum((s-1)*st for s,st in zip(real_real_shape,strides)) + 1 - if i: buffer_size = prod(st.views[i-1].shape) - real_offset + if i: buffer_size = prod(st.views[i-1].shape) - real_offset if real_shape else 1 def sort_by_strides(shape, strides): return sorted(zip(shape, strides), key=lambda k: (k[1],-k[0]), reverse=True), sorted(range(len(strides)), key=lambda k: (strides[k],-real_real_shape[k]), reverse=True) ordered_shape_strides, order = sort_by_strides(real_real_shape, strides) to_apply.extend([(MovementOps.RESHAPE, (-1,)), (MovementOps.SHRINK, ((real_offset, real_offset+buffer_size),))]) diff --git a/extra/torch_backend/test.py b/extra/torch_backend/test.py index f92be531c1..113c013c81 100644 --- a/extra/torch_backend/test.py +++ b/extra/torch_backend/test.py @@ -170,6 +170,11 @@ class TestTorchBackend(unittest.TestCase): assert torch.equal(tensor_a, tensor_b) assert not torch.equal(tensor_a, tensor_c) + def test_scalar_assign(self): + a = torch.tensor([1, 2, 3], device=device) + a[1] = 4 + np.testing.assert_equal(a.cpu().numpy(), [1, 4, 3]) + @unittest.skip("meh") def test_str(self): a = torch.ones(4, device=device)