Repository: taichi-dev/taichi (language: python)

Issue:
[lang] Vector size not checked before scalarize
This following code:

```python
import taichi as ti

ti.init()

f = ti.Vector.field(4, dtype = ti.f32, shape = ())

@ti.kernel
def update_f():
    f[None] = ti.Vector([1, 2, 3])

update_f()
print(f)
```

produces this error:

```
[Taichi] version 1.5.0, llvm 15.0.1, commit ae61418d, win, python 3.9.12
[I 02/27/23 14:08:59.905 27420] [shell.py:_shell_pop_print@23] Graphical python shell detected, using wrapped sys.stdout
c:\users\cheng cao\source\repos\taichi\python\taichi\types\ndarray_type.py:91: DeprecationWarning: The element_dim and element_shape arguments for ndarray will be deprecated in v1.5.0, use matrix dtype instead.
  warnings.warn(
[Taichi] Starting on arch=x64
[E 02/27/23 14:09:00.348 27420] [scalarize.cpp:taichi::lang::Scalarize::scalarize_store_stmt@51] Assertion failure: dest_tensor_type->get_shape() == val_tensor_type->get_shape()


Traceback (most recent call last):
  File "C:\Users\Cheng Cao\test_vec.py", line 11, in <module>
    update_f()
  File "c:\users\cheng cao\source\repos\taichi\python\taichi\lang\kernel_impl.py", line 1004, in wrapped
    return primal(*args, **kwargs)
  File "c:\users\cheng cao\source\repos\taichi\python\taichi\lang\shell.py", line 27, in new_call
    ret = old_call(*args, **kwargs)
  File "c:\users\cheng cao\source\repos\taichi\python\taichi\lang\kernel_impl.py", line 931, in __call__
    return self.runtime.compiled_functions[key](*args)
  File "c:\users\cheng cao\source\repos\taichi\python\taichi\lang\kernel_impl.py", line 835, in func__
    raise e from None
  File "c:\users\cheng cao\source\repos\taichi\python\taichi\lang\kernel_impl.py", line 832, in func__
    t_kernel(launch_ctx)
RuntimeError: [scalarize.cpp:taichi::lang::Scalarize::scalarize_store_stmt@51] Assertion failure: dest_tensor_type->get_shape() == val_tensor_type->get_shape()
```

1. This should be a compile error
2. This error should be caught in type check


Discussion / hints:

Fix patch:
diff --git a/taichi/transforms/frontend_type_check.cpp b/taichi/transforms/frontend_type_check.cpp
index 01b97c9d5a47d..be863c4433aaf 100644
--- a/taichi/transforms/frontend_type_check.cpp
+++ b/taichi/transforms/frontend_type_check.cpp
@@ -48,7 +48,20 @@ class FrontendTypeCheck : public IRVisitor {
   }
 
   void visit(FrontendAssignStmt *stmt) override {
+    auto lhs_type = stmt->lhs->ret_type;
+    auto rhs_type = stmt->rhs->ret_type;
+
+    auto error = [&]() {
+      throw TaichiTypeError(fmt::format("{}cannot assign '{}' to '{}'",
+                                        stmt->tb, rhs_type->to_string(),
+                                        lhs_type->to_string()));
+    };
+
     // No implicit cast at frontend for now
+    if (is_tensor(lhs_type) && is_tensor(rhs_type) &&
+        lhs_type.get_shape() != rhs_type.get_shape()) {
+      error();
+    }
   }
 
   void visit(FrontendIfStmt *stmt) override {


Test patch:
diff --git a/tests/python/test_type_check.py b/tests/python/test_type_check.py
index 964e24131b106..a52a60d10569b 100644
--- a/tests/python/test_type_check.py
+++ b/tests/python/test_type_check.py
@@ -80,3 +80,18 @@ def foo():
             "Only 0-dimensional numpy array can be used to initialize a scalar expression"
     ):
         foo()
+
+
+@test_utils.test(arch=ti.cpu)
+def test_assign():
+    f = ti.Vector.field(4, dtype=ti.i32, shape=())
+
+    @ti.kernel
+    def floor():
+        f[None] = ti.Vector([1, 2, 3])
+
+    with pytest.raises(
+            ti.TaichiTypeError,
+            match=
+            r"cannot assign '\[Tensor \(3\) i32\]' to '\[Tensor \(4\) i32\]'"):
+        floor()