Improve error message for unpacking a tile (#125)

Author: jansel

helion/_compiler/type_propagation.py

@@ -1600,13 +1600,11 @@ def _assign(self, lhs: ast.AST, rhs: TypeInfo) -> None:
             try:
                 elements = rhs.unpack()
             except NotImplementedError:
-                elements = [
-                    UnknownType(
-                        self.origin(),
-                        f"Failed to unpack assignment: {rhs!s}",
-                    )
-                    for _ in lhs
-                ]
+                if isinstance(rhs, UnknownType):
+                    raise exc.TypePropagationError(rhs) from None
+                if isinstance(rhs, TileIndexType):
+                    raise exc.FailedToUnpackTile from None
+                raise exc.FailedToUnpackTupleAssign(len(lhs), rhs) from None
             used_star = False
             idx = 0
             for elt in lhs:

helion/exc.py

@@ -113,6 +113,18 @@ class SpecializeArgType(BaseError):
     message = "hl.specialize() must be called on a size from an input tensor, got: {}"
 
 
+class FailedToUnpackTupleAssign(BaseError):
+    message = "Failed to unpack values in tuple assignment.  Expected a sequence of size {0}, got type: {1!s}."
+
+
+class FailedToUnpackTile(BaseError):
+    message = (
+        "Failed to unpack a tile into a tuple assignment. "
+        "Expected an sequence, but got a single tile. "
+        "Did you mix up `hl.tile(x)` and `hl.tile([x])`?"
+    )
+
+
 class AssignmentMultipleTargets(NotAllowedOnDevice):
     message = "Assignment with multiple targets (a=b=1) is not allowed inside the `hl.tile` or `hl.grid` loop."
 

test/test_errors.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from expecttest import TestCase
+import torch
+
+import helion
+from helion._testing import DEVICE
+from helion._testing import code_and_output
+import helion.language as hl
+
+
+class TestErrors(TestCase):
+    maxDiff = 16384
+
+    def test_tile_unpacking(self):
+        @helion.kernel()
+        def sum_kernel(x: torch.Tensor) -> torch.Tensor:
+            batch, seq_len, hidden = x.size()
+            out = x.new_empty(batch, hidden)
+            for tile_batch, tile_hidden in hl.tile(batch, hidden):
+                out[tile_batch, tile_hidden] = x[tile_batch, :, tile_hidden].sum(1)
+            return out
+
+        with self.assertRaises(helion.exc.FailedToUnpackTile):
+            code_and_output(sum_kernel, (torch.randn(2, 3, 4, device=DEVICE),))