gstaichi 0.1.20.dev0__cp310-cp310-macosx_15_0_arm64.whl → 0.1.25.dev0__cp310-cp310-macosx_15_0_arm64.whl

{taichi → gstaichi}/lang/impl.py

@@ -1,33 +1,31 @@
-# type: ignore
-
 import numbers
 from types import FunctionType, MethodType
 from typing import Any, Iterable, Sequence
 
 import numpy as np
 
-from taichi._lib import core as _ti_core
-from taichi._lib.core.taichi_python import (
-    DataType,
+from gstaichi._lib import core as _ti_core
+from gstaichi._lib.core.gstaichi_python import (
+    DataTypeCxx,
     Function,
     Program,
 )
-from taichi._snode.fields_builder import FieldsBuilder
-from taichi.lang._ndarray import ScalarNdarray
-from taichi.lang._ndrange import GroupedNDRange, _Ndrange
-from taichi.lang._texture import RWTextureAccessor
-from taichi.lang.any_array import AnyArray
-from taichi.lang.exception import (
-    TaichiCompilationError,
-    TaichiRuntimeError,
-    TaichiSyntaxError,
-    TaichiTypeError,
+from gstaichi._snode.fields_builder import FieldsBuilder
+from gstaichi.lang._ndarray import ScalarNdarray
+from gstaichi.lang._ndrange import GroupedNDRange, _Ndrange
+from gstaichi.lang._texture import RWTextureAccessor
+from gstaichi.lang.any_array import AnyArray
+from gstaichi.lang.exception import (
+    GsTaichiCompilationError,
+    GsTaichiRuntimeError,
+    GsTaichiSyntaxError,
+    GsTaichiTypeError,
 )
-from taichi.lang.expr import Expr, make_expr_group
-from taichi.lang.field import Field, ScalarField
-from taichi.lang.kernel_arguments import SparseMatrixProxy
-from taichi.lang.kernel_impl import Kernel
-from taichi.lang.matrix import (
+from gstaichi.lang.expr import Expr, make_expr_group
+from gstaichi.lang.field import Field, ScalarField
+from gstaichi.lang.kernel_arguments import SparseMatrixProxy
+from gstaichi.lang.kernel_impl import BoundGsTaichiCallable, GsTaichiCallable, Kernel
+from gstaichi.lang.matrix import (
     Matrix,
     MatrixField,
     MatrixNdarray,
@@ -36,7 +34,7 @@ from taichi.lang.matrix import (
     VectorNdarray,
     make_matrix,
 )
-from taichi.lang.mesh import (
+from gstaichi.lang.mesh import (
     ConvType,
     MeshElementFieldProxy,
     MeshInstance,
@@ -45,19 +43,19 @@ from taichi.lang.mesh import (
     MeshReorderedScalarFieldProxy,
     element_type_name,
 )
-from taichi.lang.simt.block import SharedArray
-from taichi.lang.snode import SNode
-from taichi.lang.struct import Struct, StructField, _IntermediateStruct
-from taichi.lang.util import (
+from gstaichi.lang.simt.block import SharedArray
+from gstaichi.lang.snode import SNode
+from gstaichi.lang.struct import Struct, StructField, _IntermediateStruct
+from gstaichi.lang.util import (
     cook_dtype,
     get_traceback,
-    is_taichi_class,
+    gstaichi_scope,
+    is_gstaichi_class,
     python_scope,
-    taichi_scope,
     warning,
 )
-from taichi.types.enums import SNodeGradType
-from taichi.types.primitive_types import (
+from gstaichi.types.enums import SNodeGradType
+from gstaichi.types.primitive_types import (
     all_types,
     f16,
     f32,
@@ -70,7 +68,7 @@ from taichi.types.primitive_types import (
 )
 
 
-@taichi_scope
+@gstaichi_scope
 def expr_init_shared_array(shape, element_type):
     compiling_callable = get_runtime().compiling_callable
     assert compiling_callable is not None
@@ -79,7 +77,7 @@ def expr_init_shared_array(shape, element_type):
     )
 
 
-@taichi_scope
+@gstaichi_scope
 def expr_init(rhs):
     compiling_callable = get_runtime().compiling_callable
     assert compiling_callable is not None
@@ -88,7 +86,7 @@ def expr_init(rhs):
             compiling_callable.ast_builder().expr_alloca(_ti_core.DebugInfo(get_runtime().get_current_src_info()))
         )
     if isinstance(rhs, Matrix) and (hasattr(rhs, "_DIM")):
-        return Matrix(*rhs.to_list(), ndim=rhs.ndim)
+        return Matrix(*rhs.to_list(), ndim=rhs.ndim)  # type: ignore
     if isinstance(rhs, Matrix):
         return make_matrix(rhs.to_list())
     if isinstance(rhs, SharedArray):
@@ -101,7 +99,7 @@ def expr_init(rhs):
         return tuple(expr_init(e) for e in rhs)
     if isinstance(rhs, dict):
         return dict((key, expr_init(val)) for key, val in rhs.items())
-    if isinstance(rhs, _ti_core.DataType):
+    if isinstance(rhs, _ti_core.DataTypeCxx):
         return rhs
     if isinstance(rhs, _ti_core.Arch):
         return rhs
@@ -120,7 +118,7 @@ def expr_init(rhs):
     )
 
 
-@taichi_scope
+@gstaichi_scope
 def expr_init_func(rhs):  # temporary solution to allow passing in fields as arguments
     if isinstance(rhs, Field):
         return rhs
@@ -130,7 +128,7 @@ def expr_init_func(rhs):  # temporary solution to allow passing in fields as arg
 def begin_frontend_struct_for(ast_builder, group, loop_range):
     if not isinstance(loop_range, (AnyArray, Field, SNode, RWTextureAccessor, _Root)):
         raise TypeError(
-            f"Cannot loop over the object {type(loop_range)} in Taichi scope. Only Taichi fields (via template) or dense arrays (via types.ndarray) are supported."
+            f"Cannot loop over the object {type(loop_range)} in GsTaichi scope. Only GsTaichi fields (via template) or dense arrays (via types.ndarray) are supported."
         )
     if group.size() != len(loop_range.shape):
         raise IndexError(
@@ -147,7 +145,7 @@ def begin_frontend_struct_for(ast_builder, group, loop_range):
 
 def begin_frontend_if(ast_builder, cond, stmt_dbg_info):
     assert ast_builder is not None
-    if is_taichi_class(cond):
+    if is_gstaichi_class(cond):
         raise ValueError(
             "The truth value of vectors/matrices is ambiguous.\n"
             "Consider using `any` or `all` when comparing vectors/matrices:\n"
@@ -158,15 +156,15 @@ def begin_frontend_if(ast_builder, cond, stmt_dbg_info):
     ast_builder.begin_frontend_if(Expr(cond).ptr, stmt_dbg_info)
 
 
-@taichi_scope
+@gstaichi_scope
 def _calc_slice(index, default_stop):
     start, stop, step = index.start or 0, index.stop or default_stop, index.step or 1
 
     def check_validity(x):
         #  TODO(mzmzm): support variable in slice
         if isinstance(x, Expr):
-            raise TaichiCompilationError(
-                "Taichi does not support variables in slice now, please use constant instead of it."
+            raise GsTaichiCompilationError(
+                "GsTaichi does not support variables in slice now, please use constant instead of it."
             )
 
     check_validity(start), check_validity(stop), check_validity(step)
@@ -190,16 +188,16 @@ def validate_subscript_index(value, index):
         validate_subscript_index(value, index.stop)
 
     if isinstance(index, int) and index < 0:
-        raise TaichiSyntaxError("Negative indices are not supported in Taichi kernels.")
+        raise GsTaichiSyntaxError("Negative indices are not supported in GsTaichi kernels.")
 
 
-@taichi_scope
+@gstaichi_scope
 def subscript(ast_builder, value, *_indices, skip_reordered=False):
     dbg_info = _ti_core.DebugInfo(get_runtime().get_current_src_info())
     compiling_callable = get_runtime().compiling_callable
     assert compiling_callable is not None
     ast_builder = compiling_callable.ast_builder()
-    # Directly evaluate in Python for non-Taichi types
+    # Directly evaluate in Python for non-GsTaichi types
     if not isinstance(
         value,
         (
@@ -237,14 +235,14 @@ def subscript(ast_builder, value, *_indices, skip_reordered=False):
     indices_expr_group = None
     if has_slice:
         if not (isinstance(value, Expr) and value.is_tensor()):
-            raise TaichiSyntaxError(f"The type {type(value)} do not support index of slice type")
+            raise GsTaichiSyntaxError(f"The type {type(value)} do not support index of slice type")
     else:
         indices_expr_group = make_expr_group(*indices)
 
     if isinstance(value, SharedArray):
         return value.subscript(*indices)
     if isinstance(value, MeshElementFieldProxy):
-        return value.subscript(*indices)
+        return value.subscript(*indices)  # type: ignore
     if isinstance(value, MeshRelationAccessProxy):
         return value.subscript(*indices)
     if isinstance(value, (MeshReorderedScalarFieldProxy, MeshReorderedMatrixFieldProxy)) and not skip_reordered:
@@ -333,7 +331,7 @@ class SrcInfoGuard:
         self.info_stack.pop()
 
 
-class PyTaichi:
+class PyGsTaichi:
     def __init__(self, kernels=None):
         self.materialized = False
         self._prog: Program | None = None
@@ -359,13 +357,13 @@ class PyTaichi:
     @property
     def prog(self) -> Program:
         if self._prog is None:
-            raise TaichiRuntimeError("_prog attribute not initialized. Maybe you forgot to call `ti.init()` first?")
+            raise GsTaichiRuntimeError("_prog attribute not initialized. Maybe you forgot to call `ti.init()` first?")
         return self._prog
 
     @property
     def current_kernel(self) -> Kernel:
         if self._current_kernel is None:
-            raise TaichiRuntimeError(
+            raise GsTaichiRuntimeError(
                 "_pr_current_kernelog attribute not initialized. Maybe you forgot to call `ti.init()` first?"
             )
         return self._current_kernel
@@ -385,7 +383,7 @@ class PyTaichi:
             if builder == _root_fb:
                 continue
 
-            raise TaichiRuntimeError(
+            raise GsTaichiRuntimeError(
                 f"Field builder {builder} is not finalized. " f"Please call finalize() on it. Traceback:\n{tb}"
             )
 
@@ -426,7 +424,7 @@ class PyTaichi:
             # if the root itself is empty), so that there is a valid struct
             # llvm::Module, if no field has been declared before the first kernel
             # invocation. Example case:
-            # https://github.com/taichi-dev/taichi/blob/27bb1dc3227d9273a79fcb318fdb06fd053068f5/tests/python/test_ad_basics.py#L260-L266
+            # https://github.com/taichi-dev/gstaichi/blob/27bb1dc3227d9273a79fcb318fdb06fd053068f5/tests/python/test_ad_basics.py#L260-L266
             return
 
         if get_runtime().prog.config().debug:
@@ -437,13 +435,6 @@ class PyTaichi:
         global _root_fb
         _root_fb = FieldsBuilder()
 
-    @staticmethod
-    def _finalize_root_fb_for_aot():
-        if _root_fb.finalized:
-            raise RuntimeError("AOT: can only finalize the root FieldsBuilder once")
-        assert isinstance(_root_fb, FieldsBuilder)
-        _root_fb._finalize_for_aot()
-
     @staticmethod
     def _get_tb(_var):
         return getattr(_var, "declaration_tb", str(_var.ptr))
@@ -531,33 +522,33 @@ class PyTaichi:
         self._prog.synchronize()
 
 
-pytaichi = PyTaichi()
+pygstaichi = PyGsTaichi()
 
 
-def get_runtime() -> PyTaichi:
-    return pytaichi
+def get_runtime() -> PyGsTaichi:
+    return pygstaichi
 
 
 def reset():
-    global pytaichi
-    old_kernels = pytaichi.kernels
-    pytaichi.clear()
-    pytaichi = PyTaichi(old_kernels)
+    global pygstaichi
+    old_kernels = pygstaichi.kernels
+    pygstaichi.clear()
+    pygstaichi = PyGsTaichi(old_kernels)
     for k in old_kernels:
         k.reset()
     _ti_core.reset_default_compile_config()
 
 
-@taichi_scope
+@gstaichi_scope
 def static_print(*args, __p=print, **kwargs):
-    """The print function in Taichi scope.
+    """The print function in GsTaichi scope.
 
     This function is called at compile time and has no runtime overhead.
     """
     __p(*args, **kwargs)
 
 
-# we don't add @taichi_scope decorator for @ti.pyfunc to work
+# we don't add @gstaichi_scope decorator for @ti.pyfunc to work
 def static_assert(cond, msg=None):
     """Throw AssertionError when `cond` is False.
 
@@ -577,7 +568,7 @@ def static_assert(cond, msg=None):
         AssertionError: the year must be a lunar year
     """
     if isinstance(cond, Expr):
-        raise TaichiTypeError("Static assert with non-static condition")
+        raise GsTaichiTypeError("Static assert with non-static condition")
     if msg is not None:
         assert cond, msg
     else:
@@ -585,7 +576,7 @@ def static_assert(cond, msg=None):
 
 
 def inside_kernel():
-    return pytaichi.inside_kernel
+    return pygstaichi.inside_kernel
 
 
 def index_nd(dim):
@@ -597,21 +588,21 @@ class _UninitializedRootFieldsBuilder:
         if item == "__qualname__":
             # For sphinx docstring extraction.
             return "_UninitializedRootFieldsBuilder"
-        raise TaichiRuntimeError("Please call init() first")
+        raise GsTaichiRuntimeError("Please call init() first")
 
 
 # `root` initialization must be delayed until after the program is
-# created. Unfortunately, `root` exists in both taichi.lang.impl module and
-# the top-level taichi module at this point; so if `root` itself is written, we
+# created. Unfortunately, `root` exists in both gstaichi.lang.impl module and
+# the top-level gstaichi module at this point; so if `root` itself is written, we
 # would have to make sure that `root` in all the modules get updated to the same
 # instance. This is an error-prone process.
 #
 # To avoid this situation, we create `root` once during the import time, and
 # never write to it. The core part, `_root_fb`, is the one whose initialization
-# gets delayed. `_root_fb` will only exist in the taichi.lang.impl module, so
+# gets delayed. `_root_fb` will only exist in the gstaichi.lang.impl module, so
 # writing to it is would result in less for maintenance cost.
 #
-# `_root_fb` will be overridden inside :func:`taichi.lang.init`.
+# `_root_fb` will be overridden inside :func:`gstaichi.lang.init`.
 _root_fb = _UninitializedRootFieldsBuilder()
 
 
@@ -626,19 +617,19 @@ class _Root:
 
     @staticmethod
     def parent(n=1):
-        """Same as :func:`taichi.SNode.parent`"""
+        """Same as :func:`gstaichi.SNode.parent`"""
         assert isinstance(_root_fb, FieldsBuilder)
         return _root_fb.root.parent(n)
 
     @staticmethod
     def _loop_range():
-        """Same as :func:`taichi.SNode.loop_range`"""
+        """Same as :func:`gstaichi.SNode.loop_range`"""
         assert isinstance(_root_fb, FieldsBuilder)
         return _root_fb.root._loop_range()
 
     @staticmethod
     def _get_children():
-        """Same as :func:`taichi.SNode.get_children`"""
+        """Same as :func:`gstaichi.SNode.get_children`"""
         assert isinstance(_root_fb, FieldsBuilder)
         return _root_fb.root._get_children()
 
@@ -650,7 +641,7 @@ class _Root:
 
     @property
     def shape(self):
-        """Same as :func:`taichi.SNode.shape`"""
+        """Same as :func:`gstaichi.SNode.shape`"""
         assert isinstance(_root_fb, FieldsBuilder)
         return _root_fb.root.shape
 
@@ -667,7 +658,7 @@ class _Root:
 
 
 root = _Root()
-"""Root of the declared Taichi :func:`~taichi.lang.impl.field`s.
+"""Root of the declared GsTaichi :func:`~gstaichi.lang.impl.field`s.
 
 See also https://docs.taichi-lang.org/docs/layout
 
@@ -702,7 +693,7 @@ def create_field_member(dtype, name, needs_grad, needs_dual):
     x.ptr = _ti_core.expr_field(x.ptr, dtype)
     x.ptr.set_name(name)
     x.ptr.set_grad_type(SNodeGradType.PRIMAL)
-    pytaichi.global_vars.append(x)
+    pygstaichi.global_vars.append(x)
 
     x_grad = None
     x_dual = None
@@ -717,13 +708,13 @@ def create_field_member(dtype, name, needs_grad, needs_dual):
         x_grad.ptr.set_grad_type(SNodeGradType.ADJOINT)
         x.ptr.set_adjoint(x_grad.ptr)
         if needs_grad:
-            pytaichi.grad_vars.append(x_grad)
+            pygstaichi.grad_vars.append(x_grad)
 
         if prog.config().debug:
             # adjoint checkbit
             x_grad_checkbit = Expr(prog.make_id_expr(""))
             dtype = u8
-            if prog.config().arch in (_ti_core.opengl, _ti_core.vulkan, _ti_core.gles):
+            if prog.config().arch == _ti_core.vulkan:
                 dtype = i32
             x_grad_checkbit.ptr = _ti_core.expr_field(x_grad_checkbit.ptr, cook_dtype(dtype))
             x_grad_checkbit.ptr.set_name(name + ".grad_checkbit")
@@ -737,9 +728,9 @@ def create_field_member(dtype, name, needs_grad, needs_dual):
         x_dual.ptr.set_grad_type(SNodeGradType.DUAL)
         x.ptr.set_dual(x_dual.ptr)
         if needs_dual:
-            pytaichi.dual_vars.append(x_dual)
+            pygstaichi.dual_vars.append(x_dual)
     elif needs_grad or needs_dual:
-        raise TaichiRuntimeError(f"{dtype} is not supported for field with `needs_grad=True` or `needs_dual=True`.")
+        raise GsTaichiRuntimeError(f"{dtype} is not supported for field with `needs_grad=True` or `needs_dual=True`.")
 
     return x, x_grad, x_dual
 
@@ -765,9 +756,9 @@ def _field(
 
     if shape is None:
         if offset is not None:
-            raise TaichiSyntaxError("shape cannot be None when offset is set")
+            raise GsTaichiSyntaxError("shape cannot be None when offset is set")
         if order is not None:
-            raise TaichiSyntaxError("shape cannot be None when order is set")
+            raise GsTaichiSyntaxError("shape cannot be None when order is set")
     else:
         if isinstance(shape, numbers.Number):
             shape = (shape,)
@@ -775,20 +766,22 @@ def _field(
             offset = (offset,)
         dim = len(shape)
         if offset is not None and dim != len(offset):
-            raise TaichiSyntaxError(f"The dimensionality of shape and offset must be the same ({dim} != {len(offset)})")
+            raise GsTaichiSyntaxError(
+                f"The dimensionality of shape and offset must be the same ({dim} != {len(offset)})"
+            )
         axis_seq = []
         shape_seq = []
         if order is not None:
             if dim != len(order):
-                raise TaichiSyntaxError(
+                raise GsTaichiSyntaxError(
                     f"The dimensionality of shape and order must be the same ({dim} != {len(order)})"
                 )
             if dim != len(set(order)):
-                raise TaichiSyntaxError("The axes in order must be different")
+                raise GsTaichiSyntaxError("The axes in order must be different")
             for ch in order:
                 axis = ord(ch) - ord("i")
                 if axis < 0 or axis >= dim:
-                    raise TaichiSyntaxError(f"Invalid axis {ch}")
+                    raise GsTaichiSyntaxError(f"Invalid axis {ch}")
                 axis_seq.append(axis)
                 shape_seq.append(shape[axis])
         else:
@@ -805,12 +798,12 @@ def _field(
 
 @python_scope
 def field(dtype, *args, **kwargs):
-    """Defines a Taichi field.
+    """Defines a GsTaichi field.
 
-    A Taichi field can be viewed as an abstract N-dimensional array, hiding away
-    the complexity of how its underlying :class:`~taichi.lang.snode.SNode` are
-    actually defined. The data in a Taichi field can be directly accessed by
-    a Taichi :func:`~taichi.lang.kernel_impl.kernel`.
+    A GsTaichi field can be viewed as an abstract N-dimensional array, hiding away
+    the complexity of how its underlying :class:`~gstaichi.lang.snode.SNode` are
+    actually defined. The data in a GsTaichi field can be directly accessed by
+    a GsTaichi :func:`~gstaichi.lang.kernel_impl.kernel`.
 
     See also https://docs.taichi-lang.org/docs/field
 
@@ -827,7 +820,7 @@ def field(dtype, *args, **kwargs):
 
     Example::
 
-        The code below shows how a Taichi field can be declared and defined::
+        The code below shows how a GsTaichi field can be declared and defined::
 
             >>> x1 = ti.field(ti.f32, shape=(16, 8))
             >>> # Equivalently
@@ -851,14 +844,14 @@ def field(dtype, *args, **kwargs):
 
 @python_scope
 def ndarray(dtype, shape, needs_grad=False):
-    """Defines a Taichi ndarray with scalar elements.
+    """Defines a GsTaichi ndarray with scalar elements.
 
     Args:
         dtype (Union[DataType, MatrixType]): Data type of each element. This can be either a scalar type like ti.f32 or a compound type like ti.types.vector(3, ti.i32).
         shape (Union[int, tuple[int]]): Shape of the ndarray.
 
     Example:
-        The code below shows how a Taichi ndarray with scalar elements can be declared and defined::
+        The code below shows how a GsTaichi ndarray with scalar elements can be declared and defined::
 
             >>> x = ti.ndarray(ti.f32, shape=(16, 8))  # ndarray of shape (16, 8), each element is ti.f32 scalar.
             >>> vec3 = ti.types.vector(3, ti.i32)
@@ -870,7 +863,7 @@ def ndarray(dtype, shape, needs_grad=False):
     if isinstance(shape, numbers.Number):
         shape = (shape,)
     if not all((isinstance(x, int) or isinstance(x, np.integer)) and x > 0 and x <= 2**31 - 1 for x in shape):
-        raise TaichiRuntimeError(f"{shape} is not a valid shape for ndarray")
+        raise GsTaichiRuntimeError(f"{shape} is not a valid shape for ndarray")
     if dtype in all_types:
         dt = cook_dtype(dtype)
         x = ScalarNdarray(dt, shape)
@@ -881,17 +874,19 @@ def ndarray(dtype, shape, needs_grad=False):
             x = MatrixNdarray(dtype.n, dtype.m, dtype.dtype, shape)
         dt = dtype.dtype
     else:
-        raise TaichiRuntimeError(f"{dtype} is not supported as ndarray element type")
+        raise GsTaichiRuntimeError(f"{dtype} is not supported as ndarray element type")
     if needs_grad:
-        assert isinstance(dt, DataType)
+        assert isinstance(dt, DataTypeCxx)
         if not _ti_core.is_real(dt):
-            raise TaichiRuntimeError(f"{dt} is not supported for ndarray with `needs_grad=True` or `needs_dual=True`.")
+            raise GsTaichiRuntimeError(
+                f"{dt} is not supported for ndarray with `needs_grad=True` or `needs_dual=True`."
+            )
         x_grad = ndarray(dtype, shape, needs_grad=False)
         x._set_grad(x_grad)
     return x
 
 
-@taichi_scope
+@gstaichi_scope
 def ti_format_list_to_content_entries(raw):
     # return a pair of [content, format]
     def entry2content(_var):
@@ -919,7 +914,7 @@ def ti_format_list_to_content_entries(raw):
                 yield _var[1:]
                 continue
             elif hasattr(_var, "__ti_repr__"):
-                res = _var.__ti_repr__()
+                res = _var.__ti_repr__()  # type: ignore
             elif isinstance(_var, (list, tuple)):
                 # If the first element is '__ti_format__', this list is the result of ti_format.
                 if len(_var) > 0 and isinstance(_var[0], str) and _var[0] == "__ti_format__":
@@ -956,7 +951,7 @@ def ti_format_list_to_content_entries(raw):
     return extract_formats(entries)
 
 
-@taichi_scope
+@gstaichi_scope
 def ti_print(*_vars, sep=" ", end="\n"):
     def add_separators(_vars):
         for i, _var in enumerate(_vars):
@@ -974,7 +969,7 @@ def ti_print(*_vars, sep=" ", end="\n"):
     )
 
 
-@taichi_scope
+@gstaichi_scope
 def ti_format(*args):
     content = args[0]
     mixed = args[1:]
@@ -997,7 +992,7 @@ def ti_format(*args):
     return res
 
 
-@taichi_scope
+@gstaichi_scope
 def ti_assert(cond, msg, extra_args, dbg_info):
     # Mostly a wrapper to help us convert from Expr (defined in Python) to
     # _ti_core.Expr (defined in C++)
@@ -1006,35 +1001,35 @@ def ti_assert(cond, msg, extra_args, dbg_info):
     compiling_callable.ast_builder().create_assert_stmt(Expr(cond).ptr, msg, extra_args, dbg_info)
 
 
-@taichi_scope
+@gstaichi_scope
 def ti_int(_var):
     if hasattr(_var, "__ti_int__"):
         return _var.__ti_int__()
     return int(_var)
 
 
-@taichi_scope
+@gstaichi_scope
 def ti_bool(_var):
     if hasattr(_var, "__ti_bool__"):
         return _var.__ti_bool__()
     return bool(_var)
 
 
-@taichi_scope
+@gstaichi_scope
 def ti_float(_var):
     if hasattr(_var, "__ti_float__"):
         return _var.__ti_float__()
     return float(_var)
 
 
-@taichi_scope
+@gstaichi_scope
 def zero(x):
     # TODO: get dtype from Expr and Matrix:
     """Returns an array of zeros with the same shape and type as the input. It's also a scalar
     if the input is a scalar.
 
     Args:
-        x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): The input.
+        x (Union[:mod:`~gstaichi.types.primitive_types`, :class:`~gstaichi.Matrix`]): The input.
 
     Returns:
         A new copy of the input but filled with zeros.
@@ -1051,13 +1046,13 @@ def zero(x):
     return x * 0
 
 
-@taichi_scope
+@gstaichi_scope
 def one(x):
     """Returns an array of ones with the same shape and type as the input. It's also a scalar
     if the input is a scalar.
 
     Args:
-        x (Union[:mod:`~taichi.types.primitive_types`, :class:`~taichi.Matrix`]): The input.
+        x (Union[:mod:`~gstaichi.types.primitive_types`, :class:`~gstaichi.Matrix`]): The input.
 
     Returns:
         A new copy of the input but filled with ones.
@@ -1080,7 +1075,7 @@ def axes(*x: int):
     Args:
         *x: A list of axes to be activated
 
-    Note that Taichi has already provided a set of commonly used axes. For example,
+    Note that GsTaichi has already provided a set of commonly used axes. For example,
     `ti.ij` is just `axes(0, 1)` under the hood.
     """
     return [_ti_core.Axis(i) for i in x]
@@ -1090,9 +1085,9 @@ Axis = _ti_core.Axis
 
 
 def static(x, *xs) -> Any:
-    """Evaluates a Taichi-scope expression at compile time.
+    """Evaluates a GsTaichi-scope expression at compile time.
 
-    `static()` is what enables the so-called metaprogramming in Taichi. It is
+    `static()` is what enables the so-called metaprogramming in GsTaichi. It is
     in many ways similar to ``constexpr`` in C++.
 
     See also https://docs.taichi-lang.org/docs/meta.
@@ -1162,12 +1157,12 @@ def static(x, *xs) -> Any:
         return x
     if isinstance(x, Field):
         return x
-    if isinstance(x, (FunctionType, MethodType)):
+    if isinstance(x, (FunctionType, MethodType, BoundGsTaichiCallable, GsTaichiCallable)):
         return x
     raise ValueError(f"Input to ti.static must be compile-time constants or global pointers, instead of {type(x)}")
 
 
-@taichi_scope
+@gstaichi_scope
 def grouped(x):
     """Groups the indices in the iterator returned by `ndrange()` into a 1-D vector.
 
@@ -1175,7 +1170,7 @@ def grouped(x):
     in one `for` loop and a single index.
 
     Args:
-        x (:func:`~taichi.ndrange`): an iterator object returned by `ti.ndrange`.
+        x (:func:`~gstaichi.ndrange`): an iterator object returned by `ti.ndrange`.
 
     Example::
         >>> # without ti.grouped
@@ -1197,7 +1192,7 @@ def stop_grad(x):
     """Stops computing gradients during back propagation.
 
     Args:
-        x (:class:`~taichi.Field`): A field.
+        x (:class:`~gstaichi.Field`): A field.
     """
     compiling_callable = get_runtime().compiling_callable
     assert compiling_callable is not None
@@ -1220,7 +1215,7 @@ def get_cuda_compute_capability():
     return _ti_core.query_int64("cuda_compute_capability")
 
 
-@taichi_scope
+@gstaichi_scope
 def mesh_relation_access(mesh, from_index, to_element_type):
     # to support ti.mesh_local and access mesh attribute as field
     if isinstance(from_index, MeshInstance):