fix: handle buffer size calculation in to_movement_ops and add scalar assignment test in torch_backend (#10464)

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),))])

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)