Fix broken tests

python/freetensor/core/expr.py

@@ -158,8 +158,9 @@ def _parse_key(self, key):
             key = (key,)
         ffiIdx = []
         if len(key) > self.ndim:
+            key_str = ', '.join(map(str, key))
             raise ffi.InvalidProgram(
-                f"Too many indices for {self.name}, expected no more than {self.ndim}"
+                f"Too many indices for {self.name}, expected no more than {self.ndim}, got [{key_str}]"
             )
         for idx, length in zip(key, self.shape()):
             if isinstance(idx, slice):

test/20.pass/test_prop_one_time_use.py

@@ -129,17 +129,17 @@ def test_used_in_many_reads_no_prop():
                     ("y", (4,), "int32", "output", "cpu")]) as (x, y):
         with ft.For("i", 0, 4) as i:
             with ft.VarDef("t", (), "int32", "cache", "cpu") as t:
-                t[i] = x[i] + 1
-                y[i] = t[i] * t[i] + t[i]
+                t[...] = x[i] + 1
+                y[i] = t[...] * t[...] + t[...]
     ast = ft.pop_ast(verbose=True)
     ast = ft.lower(ast, verbose=1)
 
     with ft.VarDef([("x", (4,), "int32", "input", "cpu"),
                     ("y", (4,), "int32", "output", "cpu")]) as (x, y):
         with ft.For("i", 0, 4) as i:
             with ft.VarDef("t", (), "int32", "cache", "cpu") as t:
-                t[i] = x[i] + 1
-                y[i] = t[i] * t[i] + t[i]
+                t[...] = x[i] + 1
+                y[i] = t[...] * t[...] + t[...]
     std = ft.pop_ast()
 
     assert std.match(ast)
@@ -150,19 +150,19 @@ def test_used_in_many_iterations_no_prop():
                     ("y", (4, 8), "int32", "output", "cpu")]) as (x, y):
         with ft.For("i", 0, 4) as i:
             with ft.VarDef("t", (), "int32", "cache", "cpu") as t:
-                t[i] = x[i] + 1
+                t[...] = x[i] + 1
                 with ft.For("j", 0, 8) as j:
-                    y[i, j] = t[i] * 2
+                    y[i, j] = t[...] * 2
     ast = ft.pop_ast(verbose=True)
     ast = ft.lower(ast, verbose=1)
 
     with ft.VarDef([("x", (4,), "int32", "input", "cpu"),
                     ("y", (4, 8), "int32", "output", "cpu")]) as (x, y):
         with ft.For("i", 0, 4) as i:
             with ft.VarDef("t", (), "int32", "cache", "cpu") as t:
-                t[i] = x[i] + 1
+                t[...] = x[i] + 1
                 with ft.For("j", 0, 8) as j:
-                    y[i, j] = t[i] * 2
+                    y[i, j] = t[...] * 2
     std = ft.pop_ast()
 
     assert std.match(ast)

test/20.pass/test_remove_cyclic_assign.py

@@ -46,7 +46,7 @@ def test_in_branch_2():
         with ft.For("i", 0, 4) as i:
             with ft.VarDef("cond", (), "bool", "cache", "cpu") as cond:
                 cond[...] = p[i] > 0
-                with ft.If(cond[i]):
+                with ft.If(cond[...]):
                     b[i] = a[i]
                     a[i] = b[i]
     ast = ft.pop_ast(verbose=True)
@@ -60,7 +60,7 @@ def test_in_branch_2():
         with ft.For("i", 0, 4) as i:
             with ft.VarDef("cond", (), "bool", "cache", "cpu") as cond:
                 cond[...] = p[i] > 0
-                with ft.If(cond[i]):
+                with ft.If(cond[...]):
                     b[i] = a[i]
     std = ft.pop_ast()
 

test/20.pass/test_remove_writes.py

@@ -783,7 +783,7 @@ def test_circular_dependence_in_parallel():
 
 
 def test_one_loop_depends_on_multiple_statements_no_remove():
-    with ft.VarDef("u", (64,), "float64", "input", "cpu") as u:
+    with ft.VarDef("u", (2, 2), "float64", "input", "cpu") as u:
         with ft.VarDef("y", (2,), "float64", "output", "cpu") as y:
             with ft.VarDef("tmp", (2,), "float64", "cache", "cpu") as tmp:
                 with ft.For("i", 0, 2) as i:
@@ -807,7 +807,7 @@ def test_one_loop_depends_on_multiple_statements_no_remove():
                        'prop_one_time_use', 'make_heap_alloc', 'float_simplify'
                    ])
 
-    with ft.VarDef("u", (64,), "float64", "input", "cpu") as u:
+    with ft.VarDef("u", (2, 2), "float64", "input", "cpu") as u:
         with ft.VarDef("y", (2,), "float64", "output", "cpu") as y:
             with ft.VarDef("tmp", (2,), "float64", "cache", "cpu") as tmp:
                 with ft.VarDef("A", (2,), "float32", "cache", "cpu") as A:

test/20.pass/test_scalar_prop_const.py

@@ -174,16 +174,16 @@ def test_prop_iter_in_expr():
     with ft.VarDef([("y1", (), "int32", "output", "cpu"),
                     ("y2", (4,), "int32", "output", "cpu")]) as (y1, y2):
         with ft.For("i", 0, 4) as i:
-            y1[i] = i + 1
-            y2[i] = y1[i]
+            y1[...] = i + 1
+            y2[i] = y1[...]
     ast = ft.pop_ast(verbose=True)
     ast = ft.lower(ast, verbose=1, skip_passes=['tensor_prop_const'])
 
     with ft.VarDef([("y1", (), "int32", "output", "cpu"),
                     ("y2", (4,), "int32", "output", "cpu")]) as (y1, y2):
         with ft.For("i", 0, 4) as i:
             with ft.If(i == 3):
-                y1[i] = 4
+                y1[...] = 4
             y2[i] = i + 1
     std = ft.pop_ast()
 

test/40.codegen/gpu/test_gpu.py

@@ -59,7 +59,7 @@ def test(x, y):
                 x16[j] = ft.cast(x[i, j], "float16")
             y16[...] = 0
             for j in range(4):
-                y16[j] += x16[j]
+                y16[...] += x16[j]
             y[i] = ft.cast(y16[...], "float32")
 
     with device: