Variable not allowed in type expression

[clouds56]

Concisely describe the proposed feature I would like to make ti.types.ndarray, tm.vec2, etc. more pylint friendly

import taichi as ti
import taichi.math as tm
from typing import TYPE_CHECKING, TypeVar, Any
img2d = ti.types.ndarray(dtype=tm.vec3)
if TYPE_CHECKING:
  from numpy import ndarray
  img2d = TypeVar("_ArrayType", bound=ndarray[Any, Any])
# this would warn about "Variable not allowed in type expression"
def ti_line(img: img2d, p0: tm.vec2, p1: tm.vec2, color: tm.vec3):
  pass

Describe the solution you'd like (if any) Have a pyi file or check if TYPE_CHECKING, make tm.vec3 a TypeVar, just like what numpy/img2d do

Additional comments