RuntimeError saying statements.inc.h:visit@19 not supported

[Wilbur-Django]

Describe the bug
Taichi a RuntimeError saying that [statements.inc.h:visit@19] Not supported on a very simple case.

To Reproduce

import taichi as ti

ti.init(default_fp=ti.f32)

@ti.data_oriented
class A:
    @ti.kernel
    def f(self) -> float:
        return 1.0


a = A()
print(a.f())  # This works
loss = ti.field(dtype=ti.f32, shape=(), needs_grad=True)
with ti.ad.Tape(loss):
    # loss[None] = 1.0  # This also works
    loss[None] = a.f() # This would cause a RuntimeError

print(loss[None])

Log/Screenshots

$ python my_sample_code.py
[Taichi] version 1.3.0, llvm 15.0.4, commit 0f25b95e, linux, python 3.9.16
[Taichi] Starting on arch=x64
1.0
[E 02/16/23 15:55:43.467 59153] [statements.inc.h:visit@19] Not supported.


Traceback (most recent call last):
  File "/home/user/my_sample_code.py", line 17, in <module>
    loss[None] = a.f() # This would cause a RuntimeError
  File "/home/user/miniconda3/envs/taichi/lib/python3.9/site-packages/taichi/ad/_ad.py", line 214, in __exit__
    self.grad()
  File "/home/user/miniconda3/envs/taichi/lib/python3.9/site-packages/taichi/ad/_ad.py", line 239, in grad
    func.grad(*args)
  File "/home/user/miniconda3/envs/taichi/lib/python3.9/site-packages/taichi/lang/kernel_impl.py", line 872, in __call__
    return self.runtime.compiled_functions[key](*args)
  File "/home/user/miniconda3/envs/taichi/lib/python3.9/site-packages/taichi/lang/kernel_impl.py", line 797, in func__
    raise e from None
  File "/home/user/miniconda3/envs/taichi/lib/python3.9/site-packages/taichi/lang/kernel_impl.py", line 794, in func__
    t_kernel(launch_ctx)
RuntimeError: [statements.inc.h:visit@19] Not supported.

Additional comments
The RuntimeError message is extremely confusing as it barely provides any information as to why the error happens. The error happens on both taichi 1.4.1 and 1.3.0.

[erizmr]

Thanks for reporting this. The RuntimeError is raised because the return statement is not properly handled by the autodiff. Have opened a PR for fixing it.

[Wilbur-Django]

Thanks for fixing it. Closing it for now.