Fix scalar tensor broadcasting in type propagation (#425)

Author: yf225

helion/_compiler/type_propagation.py

@@ -477,7 +477,8 @@ def propagate_setitem(
             lhs_rank = len(lhs_shape)
             if isinstance(value, TensorType):
                 rhs_rank = value.fake_value.ndim
-                if lhs_rank != rhs_rank:
+                # Allow scalar tensors (rank 0) to be assigned to any rank (broadcasts)
+                if rhs_rank != 0 and lhs_rank != rhs_rank:
                     raise exc.RankMismatch(
                         lhs_rank,
                         rhs_rank,

test/test_errors.py

@@ -60,24 +60,6 @@ def fn(x: torch.Tensor) -> torch.Tensor:
                 (torch.randn(4, 8, 16, device=DEVICE),),
             )
 
-    def test_rank_mismatch_assignment(self):
-        """Test that RankMismatch shows tensor shapes in assignment errors."""
-
-        @helion.kernel()
-        def fn(x: torch.Tensor) -> torch.Tensor:
-            batch, seq_len = x.size()
-            out = x.new_empty(batch, seq_len)
-            for tile_batch, tile_seq in hl.tile([batch, seq_len]):
-                scalar_val = x[tile_batch, 0].sum()  # Creates 0D tensor
-                out[tile_batch, tile_seq] = scalar_val  # 0D -> 2D assignment
-            return out
-
-        with self.assertRaisesRegex(
-            helion.exc.RankMismatch,
-            r"Expected ndim=2, but got ndim=0.*You have too few indices",
-        ):
-            code_and_output(fn, (torch.randn(4, 8, device=DEVICE),))
-
     def test_rank_mismatch_indexing(self):
         """Test that RankMismatch shows tensor shapes in indexing errors."""
 

test/test_indexing.py

@@ -608,9 +608,6 @@ def kernel(
         torch.testing.assert_close(results[4], expected_symint)
         torch.testing.assert_close(results[5], expected_symint)
 
-    @skipIfNormalMode(
-        "RankMismatch: Expected ndim=1, but got ndim=0 - LHS/RHS shape mismatch in type_propagation.py"
-    )
     def test_1d_slice_from_indexed_value(self):
         """buf[:] = zeros[i] - Assign slice from indexed value"""
 
@@ -866,9 +863,6 @@ def kernel(buf: torch.Tensor, zeros: torch.Tensor) -> torch.Tensor:
         expected = torch.zeros([N], device=DEVICE)
         torch.testing.assert_close(result, expected)
 
-    @skipIfNormalMode(
-        "RankMismatch: Expected ndim=1, but got ndim=0 - broadcasting shape mismatch"
-    )
     def test_broadcast(self):
         """Test both setter from scalar and getter for [:, i]"""