From 8926716f9327fdba26dea04b5ba1a4fd785da891 Mon Sep 17 00:00:00 2001
From: Yi Xu <xy_xuyi@foxmail.com>
Date: Fri, 3 Mar 2023 16:33:42 +0800
Subject: [PATCH] [Lang] Type check assignments between tensors (#7480)

Issue: fix #7453

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 taichi/transforms/frontend_type_check.cpp | 13 +++++++++++++
 tests/python/test_type_check.py           | 15 +++++++++++++++
 2 files changed, 28 insertions(+)

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 {
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()