关于在Taichi中如何使用类型提示(Type annotation)的疑问

综合讨论区
#1

各位好，在使用Taichi的过程中，我发现似乎无法直接在定义函数时加上类型提示。

from __future__ import annotations
import taichi as ti
ti.init(arch=ti.cpu)
@ti.func
def mul(a: ti.float16(), b: ti.float16()) -> ti.float16():
    return a*b

@ti.kernel
def foo():
    a = ti.float16(2.0)
    b = ti.float16(3.0)
    print(mul(a, b))

foo()

报错，内容是

TaichiSyntaxError: Invalid type annotation (argument 0) of Taichi function: ti.float16()

在以往的讨论中，发现大家似乎只提到了关于ndarray的情况，但按照#6中的提示，将ti.float16()改为ti.template()之后，也没有正确运行。

报错是

TaichiSyntaxError: Invalid type annotation (argument 0) of Taichi function: ti.template()

我想知道，该如何正确地进行类型提示？谢谢。

#2

改成这样就可以了

def mul(a: ti.types.f16, b: ti.types.f16) -> ti.types.f16:
1 个赞