Assertion failure: stmt->offset->ret_type.get_element_type()->is_primitive( PrimitiveTypeID::i32)

[emailweixu]

Describe the bug

It seems that taichi complains about the return type.

To Reproduce

Please post a minimal sample code to reproduce the bug.
The developer team will put a higher priority on bugs that can be reproduced within 20 lines of code. If you want a prompt reply, please keep the sample code short and representative.

import taichi as ti
ti.init(arch=ti.cpu, default_fp=ti.f64, default_ip=ti.i64)

vec4 = ti.math.vec3
mat32 = ti.types.matrix(3, 2, float)

@ti.kernel
def get_element(mat: mat32, i: ti.i32) -> float:
        return mat[0, i]
        # error goes away if change to the following. But this would be very inconvenient.
        # return mat[ti.i32(0), i]

m=mat32([[1,2], [3,4], [5,6]])
get_element(m, 0)

Log/Screenshots
Please post the full log of the program (instead of just a few lines around the error message, unless the log is > 1000 lines). This will help us diagnose what's happening. For example:

$ python my_sample_code.py
[Taichi] version 1.7.3, llvm 15.0.4, commit 5ec301be, linux, python 3.10.12
[Taichi] Starting on arch=x64
[E 02/21/25 22:42:33.872 225480] [type_check.cpp:visit@161] Assertion failure: stmt->offset->ret_type.get_element_type()->is_primitive( PrimitiveTypeID::i32)


Traceback (most recent call last):
  File "/data/weixu/code/hosim/notebooks/test.py", line 12, in <module>
    get_element(m, 0)
  File "/home/weixu/venvs/working/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 1113, in wrapped
    return primal(*args, **kwargs)
  File "/home/weixu/venvs/working/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 1045, in __call__
    return self.launch_kernel(kernel_cpp, *args)
  File "/home/weixu/venvs/working/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 976, in launch_kernel
    raise e from None
  File "/home/weixu/venvs/working/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 969, in launch_kernel
    compiled_kernel_data = prog.compile_kernel(prog.config(), prog.get_device_caps(), t_kernel)
RuntimeError: [type_check.cpp:visit@161] Assertion failure: stmt->offset->ret_type.get_element_type()->is_primitive( PrimitiveTypeID::i32)

Additional comments
If possible, please also consider attaching the output of command ti diagnose. This produces the detailed environment information and hopefully helps us diagnose faster.

ti diagnose output

[Taichi] version 1.7.3, llvm 15.0.4, commit 5ec301be, linux, python 3.10.12

*******************************************
**      Taichi Programming Language      **
*******************************************

Docs:   https://docs.taichi-lang.org/
GitHub: https://github.com/taichi-dev/taichi/
Forum:  https://forum.taichi.graphics/

Taichi system diagnose:

python: 3.10.12 (main, Mar 22 2024, 16:50:05) [GCC 11.4.0]
system: linux
executable: /home/weixu/venvs/working/bin/python
platform: Linux-6.5.0-41-generic-x86_64-with-glibc2.35
architecture: 64bit ELF
uname: uname_result(system='Linux', node='GAIL-X', release='6.5.0-41-generic', version='#41~22.04.2-Ubuntu SMP PREEMPT_DYNAMIC Mon Jun  3 11:32:55 UTC 2', machine='x86_64')
locale: en_US.UTF-8
PATH: /home/weixu/venvs/working/bin:/data/weixu/bin:/home/weixu/.local/bin:/home/weixu/miniconda3/condabin:/data/weixu/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin
PYTHONPATH: ['/home/weixu/venvs/working/bin', '/usr/lib/python310.zip', '/usr/lib/python3.10', '/usr/lib/python3.10/lib-dynload', '/home/weixu/venvs/working/lib/python3.10/site-packages', '__editable__.omni_isaac_lab-0.22.4.finder.__path_hook__', '__editable__.omni_isaac_lab_assets-0.1.3.finder.__path_hook__', '__editable__.omni_isaac_lab_tasks-0.10.0.finder.__path_hook__', '/data/weixu/code/Hobot', '/data/weixu/code/hobot_submission', '/data/weixu/code/S5']

`lsb_release` not available: [Errno 2] No such file or directory: 'lsb_release'


import: <module 'taichi' from '/home/weixu/venvs/working/lib/python3.10/site-packages/taichi/__init__.py'>

cpu: True
metal: False
opengl: True
cuda: True
vulkan: True

`glewinfo` not available: [Errno 2] No such file or directory: 'glewinfo'

Fri Feb 21 22:47:34 2025       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce RTX 3090        On  | 00000000:17:00.0 Off |                  N/A |
|  0%   56C    P2              74W / 420W |   1619MiB / 24576MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce RTX 3090        On  | 00000000:65:00.0 Off |                  N/A |
|  0%   46C    P8              42W / 420W |     14MiB / 24576MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      1652      G   /usr/lib/xorg/Xorg                          167MiB |
|    0   N/A  N/A      2078      G   /usr/bin/gnome-shell                         11MiB |
|    0   N/A  N/A     40111      C   /home/weixu/venvs/working/bin/python        350MiB |
|    0   N/A  N/A     44541      C   /home/weixu/venvs/working/bin/python        350MiB |
|    0   N/A  N/A    211010      C   /home/weixu/venvs/working/bin/python        382MiB |
|    0   N/A  N/A    216417      C   /home/weixu/venvs/working/bin/python        330MiB |
|    1   N/A  N/A      1652      G   /usr/lib/xorg/Xorg                            4MiB |
+---------------------------------------------------------------------------------------+

[Taichi] version 1.7.3, llvm 15.0.4, commit 5ec301be, linux, python 3.10.12

[Taichi] version 1.7.3, llvm 15.0.4, commit 5ec301be, linux, python 3.10.12
[Taichi] Starting on arch=x64

[Taichi] version 1.7.3, llvm 15.0.4, commit 5ec301be, linux, python 3.10.12
[Taichi] Starting on arch=opengl

[Taichi] version 1.7.3, llvm 15.0.4, commit 5ec301be, linux, python 3.10.12
[Taichi] Starting on arch=cuda

[Taichi] version 1.7.3, llvm 15.0.4, commit 5ec301be, linux, python 3.10.12

*******************************************
**      Taichi Programming Language      **
*******************************************

Docs:   https://docs.taichi-lang.org/
GitHub: https://github.com/taichi-dev/taichi/
Forum:  https://forum.taichi.graphics/

                                   TAICHI EXAMPLES                                    
 ──────────────────────────────────────────────────────────────────────────────────── 
  0: ad_gravity               25: karman_vortex_street    50: patterns                
  1: circle_packing_image     26: keyboard                51: pbf2d                   
  2: comet                    27: laplace                 52: physarum                
  3: cornell_box              28: laplace_equation        53: poisson_disk_sampling   
  4: diff_sph                 29: mandelbrot_zoom         54: print_offset            
  5: differential_evolution   30: marching_squares        55: rasterizer              
  6: euler                    31: mass_spring_3d_ggui     56: regression              
  7: eulerfluid2d             32: mass_spring_game        57: sdf_renderer            
  8: explicit_activation      33: mass_spring_game_ggui   58: simple_derivative       
  9: export_mesh              34: mciso_advanced          59: simple_texture          
  10: export_ply              35: mgpcg                   60: simple_uv               
  11: export_videos           36: mgpcg_advanced          61: snow_phaseField         
  12: fem128                  37: minimal                 62: stable_fluid            
  13: fem128_ggui             38: minimization            63: stable_fluid_ggui       
  14: fem99                   39: mpm128                  64: stable_fluid_graph      
  15: fractal                 40: mpm128_ggui             65: taichi_bitmasked        
  16: fractal3d_ggui          41: mpm3d                   66: taichi_dynamic          
  17: fullscreen              42: mpm3d_ggui              67: taichi_logo             
  18: game_of_life            43: mpm88                   68: taichi_ngp              
  19: gui_image_io            44: mpm88_graph             69: taichi_sparse           
  20: gui_widgets             45: mpm99                   70: texture_graph           
  21: implicit_fem            46: mpm_lagrangian_forces   71: tutorial                
  22: implicit_mass_spring    47: nbody                   72: two_stream_instability  
  23: initial_value_problem   48: odop_solar              73: vortex_rings            
  24: jacobian                49: oit_renderer            74: waterwave               
 ──────────────────────────────────────────────────────────────────────────────────── 
42
Running example minimal ...
[Taichi] Starting on arch=x64
42.0
>>> Running time: 0.14s

Consider attaching this log when maintainers ask about system information.
>>> Running time: 4.58s