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

Issue:
Cannot pass ndarray to a taichi function
<!-- We've collected some common issue solutions in https://docs.taichi-lang.org/docs/install. Make sure you've check them out first. Hopefully they could address your problem. -->

**Describe the bug**
Cannot pass ndarray to a taichi function

**To Reproduce**
```py
import taichi as ti

ti.init()


@ti.func
def bar(weight: ti.types.ndarray(ti.f32, ndim=3)):
    pass


@ti.kernel
def foo(weight: ti.types.ndarray(ti.f32, ndim=3)):
    bar(weight)


weight = ti.ndarray(dtype=ti.f32, shape=(2, 2, 2))
foo(weight)
```

**Log/Screenshots**

```
$ python my_sample_code.py
[Taichi] version 1.7.0, llvm 15.0.4, commit 8ca9a9e2, linux, python 3.10.0
[Taichi] Starting on arch=x64
Traceback (most recent call last):
  File "/home/lin/test/ndarray_func.py", line 17, in <module>
    foo(weight)
  File "/home/lin/taichi/python/taichi/lang/kernel_impl.py", line 944, in wrapped
    raise type(e)("\n" + str(e)) from None
taichi.lang.exception.TaichiCompilationError: 
File "/home/lin/test/ndarray_func.py", line 13, in foo:
    bar(weight)
    ^^^^^^^^^^^
File "/home/lin/test/ndarray_func.py", line 7, in bar:
def bar(weight: ti.types.ndarray(ti.f32, ndim=3)):
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Traceback (most recent call last):
  File "/home/lin/taichi/python/taichi/lang/ast/ast_transformer_utils.py", line 27, in __call__
    return method(ctx, node)
  File "/home/lin/taichi/python/taichi/lang/ast/ast_transformer.py", line 694, in build_FunctionDef
    ctx.func.arguments[i].annotation.check_matched(data.get_type())
  File "/home/lin/taichi/python/taichi/types/ndarray_type.py", line 96, in check_matched
    raise TypeError(
TypeError: Expect element type f32 for Ndarray, but get *f32
```



Discussion / hints:

Fix patch:
diff --git a/python/taichi/lang/any_array.py b/python/taichi/lang/any_array.py
index c74e774af8fad..86c6f8b44869d 100644
--- a/python/taichi/lang/any_array.py
+++ b/python/taichi/lang/any_array.py
@@ -33,7 +33,7 @@ def layout(self):
 
     def get_type(self):
         return NdarrayTypeMetadata(
-            self.ptr.get_ret_type(), None, _ti_core.get_external_tensor_needs_grad(self.ptr)
+            self.ptr.get_ret_type().ptr_removed(), None, _ti_core.get_external_tensor_needs_grad(self.ptr)
         )  # AnyArray can take any shape
 
     @property
diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
index 003bfd757e376..270c76117812a 100644
--- a/taichi/python/export_lang.cpp
+++ b/taichi/python/export_lang.cpp
@@ -111,6 +111,7 @@ void export_lang(py::module &m) {
       .def("__str__", &DataType::to_string)
       .def("shape", &DataType::get_shape)
       .def("element_type", &DataType::get_element_type)
+      .def("ptr_removed", &DataType::ptr_removed)
       .def(
           "get_ptr", [](DataType *dtype) -> Type * { return *dtype; },
           py::return_value_policy::reference)


Test patch:
diff --git a/tests/python/test_ndarray.py b/tests/python/test_ndarray.py
index 8f8586add8fd4..2429856e93093 100644
--- a/tests/python/test_ndarray.py
+++ b/tests/python/test_ndarray.py
@@ -1000,3 +1000,17 @@ def test(x: ti.types.ndarray(dtype=ti.types.vector())):
     z = ti.ndarray(ti.math.mat2, (3))
     with pytest.raises(ValueError, match=r"Invalid argument into ti.types.ndarray\(\)"):
         test(z)
+
+
+@test_utils.test(arch=supported_archs_taichi_ndarray)
+def test_pass_ndarray_to_func():
+    @ti.func
+    def bar(weight: ti.types.ndarray(ti.f32, ndim=3)):
+        pass
+
+    @ti.kernel
+    def foo(weight: ti.types.ndarray(ti.f32, ndim=3)):
+        bar(weight)
+
+    weight = ti.ndarray(dtype=ti.f32, shape=(2, 2, 2))
+    foo(weight)