gstaichi 0.1.23.dev0__cp310-cp310-win_amd64.whl → 1.0.1__cp310-cp310-win_amd64.whl

taichi/aot/module.py

@@ -1,253 +0,0 @@
-# type: ignore
-
-import datetime
-import os
-import warnings
-from contextlib import contextmanager
-from glob import glob
-from pathlib import Path, PurePosixPath
-from shutil import rmtree
-from tempfile import mkdtemp
-from zipfile import ZipFile
-
-import taichi
-from taichi.aot.utils import produce_injected_args, produce_injected_args_from_template
-from taichi.lang import impl, kernel_impl
-from taichi.lang.field import ScalarField
-from taichi.lang.matrix import MatrixField
-from taichi.types.annotations import template
-
-
-class KernelTemplate:
-    def __init__(self, kernel_fn, aot_module):
-        self._kernel_fn = kernel_fn
-        self._aot_module = aot_module
-
-    @staticmethod
-    def keygen(v, key_p, fields):
-        if isinstance(v, (int, float, bool)):
-            key_p += "=" + str(v) + ","
-            return key_p
-        for ky, val in fields:
-            if val is v:
-                key_p += "=" + ky + ","
-                return key_p
-        raise RuntimeError(
-            "Arg type must be of type int/float/boolean" f" or taichi field. Type {str(type(v))}" " is not supported"
-        )
-
-    def instantiate(self, **kwargs):
-        name = self._kernel_fn.__name__
-        kernel = self._kernel_fn._primal
-        assert isinstance(kernel, kernel_impl.Kernel)
-        injected_args = []
-        key_p = ""
-        anno_index = 0
-        template_args = {}
-
-        for index, (key, value) in enumerate(kwargs.items()):
-            template_args[index] = (key, value)
-
-        for arg in kernel.arguments:
-            if arg.annotation == template or isinstance(arg.annotation, template):
-                (k, v) = template_args[anno_index]
-                key_p += k
-                key_p = self.keygen(v, key_p, self._aot_module._fields.items())
-                injected_args.append(v)
-                anno_index += 1
-            else:
-                injected_args.append(0)
-        kernel.ensure_compiled(*injected_args)
-        self._aot_module._aot_builder.add_kernel_template(name, key_p, kernel.kernel_cpp)
-
-        # kernel AOT
-        self._aot_module._kernels.append(kernel)
-
-
-class Module:
-    """An AOT module to save and load Taichi kernels.
-
-    This module serializes the Taichi kernels for a specific arch. The
-    serialized module can later be loaded to run on that backend, without the
-    Python environment.
-
-    Example:
-      Usage::
-
-        m = ti.aot.Module(ti.metal)
-        m.add_kernel(foo)
-        m.add_kernel(bar)
-
-        m.save('/path/to/module')
-
-        # Now the module file '/path/to/module' contains the Metal kernels
-        # for running ``foo`` and ``bar``.
-    """
-
-    def __init__(self, arch=None, caps=None):
-        """Creates a new AOT module instance
-
-        Args:
-          arch: Target backend architecture. Default to the one initialized in :func:`~taichi.lang.init` if not specified.
-          caps (List[str]): Enabled device capabilities.
-        """
-        if caps is None:
-            caps = []
-        curr_arch = impl.current_cfg().arch
-        if arch is None:
-            arch = curr_arch
-        elif arch != curr_arch:
-            # TODO: we'll support this eventually but not yet...
-            warnings.warn(
-                f"AOT compilation to a different arch than the current one is not yet supported, switching to {curr_arch}"
-            )
-            arch = curr_arch
-
-        self._arch = arch
-        self._kernels = []
-        self._fields = {}
-        rtm = impl.get_runtime()
-        rtm._finalize_root_fb_for_aot()
-        self._aot_builder = rtm.prog.make_aot_module_builder(arch, caps)
-        self._content = []
-
-    def add_field(self, name, field):
-        """Add a taichi field to the AOT module.
-
-        Args:
-          name: name of taichi field
-          field: taichi field
-
-        Example::
-
-            >>> a = ti.field(ti.f32, shape=(4,4))
-            >>> b = ti.field("something")
-            >>>
-            >>> m.add_field(a)
-            >>> m.add_field(b)
-            >>>
-            >>> # Must add in sequence
-        """
-        is_scalar = True
-        self._fields[name] = field
-        column_num = 1
-        row_num = 1
-        if isinstance(field, MatrixField):
-            is_scalar = False
-            row_num = field.m
-            column_num = field.n
-        else:
-            assert isinstance(field, ScalarField)
-        self._aot_builder.add_field(
-            name,
-            field.snode.ptr,
-            is_scalar,
-            field.dtype,
-            field.snode.shape,
-            row_num,
-            column_num,
-        )
-
-    def add_kernel(self, kernel_fn, template_args=None, name=None):
-        """Add a taichi kernel to the AOT module.
-
-        Args:
-          kernel_fn (Function): the function decorated by taichi `kernel`.
-          template_args (Dict[str, Any]): a dict where key is the template
-            parameter name, and value is the instantiating arg. Note that this
-            works for both :class:`~taichi.types.template` and for
-            `:class:`~taichi.types.ndarray`.
-          name (str): Name to identify this kernel in the module. If not
-            provided, uses the built-in ``__name__`` attribute of `kernel_fn`.
-
-        """
-        kernel_name = name or kernel_fn.__name__
-        kernel = kernel_fn._primal
-        assert isinstance(kernel, kernel_impl.Kernel)
-        if template_args is not None:
-            injected_args = produce_injected_args_from_template(kernel, template_args)
-        else:
-            injected_args = produce_injected_args(kernel)
-        kernel.ensure_compiled(*injected_args)
-        self._aot_builder.add(kernel_name, kernel.kernel_cpp)
-
-        # kernel AOT
-        self._kernels.append(kernel)
-
-        self._content += ["kernel:" + kernel_name]
-
-    def add_graph(self, name, graph):
-        self._aot_builder.add_graph(name, graph._compiled_graph)
-        self._content += ["cgraph:" + name]
-
-    @contextmanager
-    def add_kernel_template(self, kernel_fn):
-        """Add a taichi kernel (with template parameters) to the AOT module.
-
-        Args:
-          kernel_fn (Function): the function decorated by taichi `kernel`.
-
-        Example::
-
-            >>> @ti.kernel
-            >>> def bar_tmpl(a: ti.template()):
-            >>>   x = a
-            >>>   # or y = a
-            >>>   # do something with `x` or `y`
-            >>>
-            >>> m = ti.aot.Module(arch)
-            >>> with m.add_kernel_template(bar_tmpl) as kt:
-            >>>   kt.instantiate(a=x)
-            >>>   kt.instantiate(a=y)
-            >>>
-            >>> @ti.kernel
-            >>> def bar_tmpl_multiple_args(a: ti.template(), b: ti.template())
-            >>>   x = a
-            >>>   y = b
-            >>>   # do something with `x` and `y`
-            >>>
-            >>> with m.add_kernel_template(bar_tmpl) as kt:
-            >>>   kt.instantiate(a=x, b=y)
-
-        TODO:
-          * Support external array
-        """
-        kt = KernelTemplate(kernel_fn, self)
-        yield kt
-
-    def save(self, filepath):
-        """
-        Args:
-          filepath (str): path to a folder to store aot files.
-        """
-        filepath = str(PurePosixPath(Path(filepath)))
-        self._aot_builder.dump(filepath, "")
-        with open(f"{filepath}/__content__", "w") as f:
-            f.write("\n".join(self._content))
-        with open(f"{filepath}/__version__", "w") as f:
-            f.write(".".join(str(x) for x in taichi.__version__))
-
-    def archive(self, filepath: str):
-        """
-        Args:
-          filepath (str): path to the stored archive of aot artifacts, MUST
-            end with `.tcm`.
-        """
-        assert filepath.endswith(".tcm"), "AOT module artifact archive must ends with .tcm"
-        tcm_path = Path(filepath).absolute()
-        assert tcm_path.parent.exists(), "Output directory doesn't exist"
-
-        temp_dir = mkdtemp(prefix="tcm_")
-        # Save first as usual.
-        self.save(temp_dir)
-
-        fixed_time = datetime.datetime(2000, 12, 1).timestamp()
-
-        # Package all artifacts into a zip archive and attach contend data.
-        with ZipFile(tcm_path, "w") as z:
-            for path in glob(f"{temp_dir}/*", recursive=True):
-                os.utime(path, (fixed_time, fixed_time))
-                z.write(path, Path.relative_to(Path(path), temp_dir))
-
-        # Remove cached files
-        rmtree(temp_dir)

taichi/aot/utils.py

@@ -1,151 +0,0 @@
-# type: ignore
-
-from typing import Any
-
-from taichi.lang._ndarray import ScalarNdarray
-from taichi.lang._texture import Texture
-from taichi.lang.exception import TaichiCompilationError
-from taichi.lang.matrix import (
-    Matrix,
-    MatrixNdarray,
-    MatrixType,
-    VectorNdarray,
-    VectorType,
-)
-from taichi.lang.util import cook_dtype
-from taichi.types.annotations import template
-from taichi.types.enums import Format
-from taichi.types.ndarray_type import NdarrayType
-from taichi.types.texture_type import RWTextureType, TextureType
-
-template_types = (NdarrayType, TextureType, template)
-
-
-def check_type_match(lhs, rhs):
-    if isinstance(lhs, MatrixType) and isinstance(rhs, MatrixType):
-        return lhs.n == rhs.n and lhs.m == rhs.m and (lhs.dtype == rhs.dtype or lhs.dtype is None or rhs.dtype is None)
-    if isinstance(lhs, MatrixType) or isinstance(rhs, MatrixType):
-        return False
-
-    return cook_dtype(lhs) == cook_dtype(rhs)
-
-
-def produce_injected_args_from_template(kernel, template_args):
-    injected_args = []
-    num_template_args = len([arg.annotation for arg in kernel.arguments if isinstance(arg.annotation, template_types)])
-    assert num_template_args == len(
-        template_args
-    ), f"Need {num_template_args} inputs to instantiate the template parameters, got {len(template_args)}"
-    for arg in kernel.arguments:
-        anno = arg.annotation
-        if isinstance(anno, template_types):
-            injected_args.append(template_args[arg.name])
-        elif isinstance(anno, RWTextureType):
-            texture_shape = (2,) * anno.num_dimensions
-            fmt = anno.fmt
-            injected_args.append(Texture(fmt, texture_shape))
-        else:
-            injected_args.append(0)
-    return injected_args
-
-
-def produce_injected_args(kernel, symbolic_args=None):
-    injected_args = []
-    for i, arg in enumerate(kernel.arguments):
-        anno = arg.annotation
-        if isinstance(anno, NdarrayType):
-            if symbolic_args is not None:
-                # TODO: reconstruct dtype to be TensorType from taichi_core instead of the Python ones
-                element_dim = len(symbolic_args[i].element_shape)
-                if element_dim == 0 or symbolic_args[i].element_shape == (1,):
-                    dtype = symbolic_args[i].dtype()
-                elif element_dim == 1:
-                    dtype = VectorType(symbolic_args[i].element_shape[0], symbolic_args[i].dtype())
-                elif element_dim == 2:
-                    dtype = MatrixType(
-                        symbolic_args[i].element_shape[0],
-                        symbolic_args[i].element_shape[1],
-                        2,
-                        symbolic_args[i].dtype(),
-                    )
-                else:
-                    raise TaichiCompilationError("Not supported")
-                ndim = symbolic_args[i].field_dim
-            else:
-                ndim = anno.ndim
-                dtype = anno.dtype
-
-            if anno.ndim is not None and ndim != anno.ndim:
-                raise TaichiCompilationError(
-                    f"{ndim} from Arg {arg.name} doesn't match kernel's annotated ndim={anno.ndim}"
-                )
-
-            if anno.dtype is not None and not check_type_match(dtype, anno.dtype):
-                raise TaichiCompilationError(
-                    f" Arg {arg.name}'s dtype {dtype.to_string()} doesn't match kernel's annotated dtype={anno.dtype.to_string()}"
-                )
-
-            if isinstance(dtype, VectorType):
-                injected_args.append(VectorNdarray(dtype.n, dtype=dtype.dtype, shape=(2,) * ndim))
-            elif isinstance(dtype, MatrixType):
-                injected_args.append(MatrixNdarray(dtype.n, dtype.m, dtype=dtype.dtype, shape=(2,) * ndim))
-            else:
-                injected_args.append(ScalarNdarray(dtype, (2,) * ndim))
-        elif isinstance(anno, RWTextureType):
-            texture_shape = (2,) * anno.num_dimensions
-            fmt = anno.fmt
-            injected_args.append(Texture(fmt, texture_shape))
-        elif isinstance(anno, TextureType):
-            texture_shape = (2,) * anno.num_dimensions
-            injected_args.append(Texture(Format.rgba8, texture_shape))
-        elif isinstance(anno, MatrixType):
-            if symbolic_args is not None:
-                symbolic_mat_n = symbolic_args[i].element_shape[0]
-                symbolic_mat_m = symbolic_args[i].element_shape[1]
-
-                if symbolic_mat_m != anno.m or symbolic_mat_n != anno.n:
-                    raise RuntimeError(
-                        f"Matrix dimension mismatch, expected ({anno.n}, {anno.m}) "
-                        f"but dispatched shape ({symbolic_mat_n}, {symbolic_mat_m})."
-                    )
-            injected_args.append(Matrix([0] * anno.n * anno.m, dt=anno.dtype))
-        else:
-            if symbolic_args is not None:
-                dtype = symbolic_args[i].dtype()
-            else:
-                dtype = anno
-
-            if not check_type_match(dtype, anno):
-                raise TaichiCompilationError(
-                    f" Arg {arg.name}'s dtype {dtype.to_string()} doesn't match kernel's annotated dtype={anno.to_string()}"
-                )
-            # For primitive types, we can just inject a dummy value.
-            injected_args.append(0)
-    return injected_args
-
-
-def json_data_model(f):
-    """
-    Decorates a JSON data model. A JSON data model MUST NOT have any member
-    functions and it MUST be constructible from a JSON object.
-
-    This is merely a marker.
-    """
-    f._is_json_data_model = True
-    return f
-
-
-def is_json_data_model(cls) -> bool:
-    return hasattr(cls, "_is_json_data_model")
-
-
-def dump_json_data_model(x: object) -> Any:
-    if isinstance(x, (int, float, str, bool, type(None))):
-        return x
-    if isinstance(x, (list, tuple)):
-        return [dump_json_data_model(e) for e in x]
-    if isinstance(x, dict):
-        return {k: dump_json_data_model(v) for k, v in x.items()}
-    if is_json_data_model(x):
-        return {k: dump_json_data_model(v) for k, v in x.__dict__.items() if k != "_is_json_data_model"}
-    return x

taichi/graph/__init__.py

@@ -1,3 +0,0 @@
-# type: ignore
-
-from ._graph import *

taichi/graph/_graph.py

@@ -1,292 +0,0 @@
-# type: ignore
-
-import warnings
-from typing import Any, Dict, List
-
-from taichi._lib import core as _ti_core
-from taichi.aot.utils import produce_injected_args
-from taichi.lang import impl, kernel_impl
-from taichi.lang._ndarray import Ndarray
-from taichi.lang._texture import Texture
-from taichi.lang.exception import TaichiRuntimeError
-from taichi.lang.matrix import Matrix, MatrixType
-from taichi.types import enums
-from taichi.types.texture_type import FORMAT2TY_CH, TY_CH2FORMAT
-
-ArgKind = _ti_core.ArgKind
-
-
-def gen_cpp_kernel(kernel_fn, args):
-    kernel = kernel_fn._primal
-    assert isinstance(kernel, kernel_impl.Kernel)
-    injected_args = produce_injected_args(kernel, symbolic_args=args)
-    key = kernel.ensure_compiled(*injected_args)
-    return kernel.compiled_kernels[key]
-
-
-def flatten_args(args):
-    unzipped_args = []
-    # Tuple for matrix args
-    # FIXME remove this when native Matrix type is ready
-    for arg in args:
-        if isinstance(arg, list):
-            for sublist in arg:
-                unzipped_args.extend(sublist)
-        else:
-            unzipped_args.append(arg)
-    return unzipped_args
-
-
-class Sequential:
-    def __init__(self, seq):
-        self.seq_ = seq
-
-    def dispatch(self, kernel_fn, *args):
-        kernel_cpp = gen_cpp_kernel(kernel_fn, args)
-        unzipped_args = flatten_args(args)
-        self.seq_.dispatch(kernel_cpp, unzipped_args)
-
-
-class GraphBuilder:
-    def __init__(self):
-        self._graph_builder = _ti_core.GraphBuilder()
-
-    def dispatch(self, kernel_fn, *args):
-        kernel_cpp = gen_cpp_kernel(kernel_fn, args)
-        unzipped_args = flatten_args(args)
-        self._graph_builder.dispatch(kernel_cpp, unzipped_args)
-
-    def create_sequential(self):
-        return Sequential(self._graph_builder.create_sequential())
-
-    def append(self, node):
-        # TODO: support appending dispatch node as well.
-        assert isinstance(node, Sequential)
-        self._graph_builder.seq().append(node.seq_)
-
-    def compile(self):
-        return Graph(self._graph_builder.compile())
-
-
-class Graph:
-    def __init__(self, compiled_graph) -> None:
-        self._compiled_graph = compiled_graph
-
-    def run(self, args):
-        # Support native python numerical types (int, float), Ndarray.
-        # Taichi Matrix types are flattened into (int, float) arrays.
-        # TODO diminish the flatten behavior when Matrix becomes a Taichi native type.
-        flattened = {}
-        for k, v in args.items():
-            if isinstance(v, Ndarray):
-                flattened[k] = v.arr
-            elif isinstance(v, Texture):
-                flattened[k] = v.tex
-            elif isinstance(v, Matrix):
-                flattened[k] = v.entries
-            elif isinstance(v, (int, float)):
-                flattened[k] = v
-            else:
-                raise TaichiRuntimeError(
-                    f"Only python int, float, ti.Matrix and ti.Ndarray are supported as runtime arguments but got {type(v)}"
-                )
-        self._compiled_graph.jit_run(impl.get_runtime().prog.config(), flattened)
-
-
-def _deprecate_arg_args(kwargs: Dict[str, Any]):
-    if "field_dim" in kwargs:
-        warnings.warn(
-            "The field_dim argument for ndarray will be deprecated in v1.6.0, use ndim instead.",
-            DeprecationWarning,
-        )
-        if "ndim" in kwargs:
-            raise TaichiRuntimeError(
-                "field_dim is deprecated, please do not specify field_dim and ndim at the same time."
-            )
-        kwargs["ndim"] = kwargs["field_dim"]
-        del kwargs["field_dim"]
-    tag = kwargs["tag"]
-
-    if tag == ArgKind.SCALAR:
-        if "element_shape" in kwargs:
-            raise TaichiRuntimeError(
-                "The element_shape argument for scalar is deprecated in v1.6.0, and is removed in v1.7.0. "
-                "Please remove them."
-            )
-
-    if tag == ArgKind.NDARRAY:
-        if "element_shape" not in kwargs:
-            if "dtype" in kwargs:
-                dtype = kwargs["dtype"]
-                if isinstance(dtype, MatrixType):
-                    kwargs["dtype"] = dtype.dtype
-                    kwargs["element_shape"] = dtype.get_shape()
-                else:
-                    kwargs["element_shape"] = ()
-        else:
-            raise TaichiRuntimeError(
-                "The element_shape argument for ndarray is deprecated in v1.6.0, and it is removed in v1.7.0. "
-                "Please use vector or matrix data type instead."
-            )
-
-    if tag == ArgKind.RWTEXTURE or tag == ArgKind.TEXTURE:
-        if "dtype" in kwargs:
-            warnings.warn(
-                "The dtype argument for texture will be deprecated in v1.6.0, use format instead.",
-                DeprecationWarning,
-            )
-            del kwargs["dtype"]
-
-        if "shape" in kwargs:
-            raise TaichiRuntimeError(
-                "The shape argument for texture is deprecated in v1.6.0, and it is removed in v1.7.0. "
-                "Please use ndim instead. (Note that you no longer need the exact texture size.)"
-            )
-
-        if "channel_format" in kwargs or "num_channels" in kwargs:
-            if "fmt" in kwargs:
-                raise TaichiRuntimeError(
-                    "channel_format and num_channels are deprecated, please do not specify channel_format/num_channels and fmt at the same time."
-                )
-            if tag == ArgKind.RWTEXTURE:
-                fmt = TY_CH2FORMAT[(kwargs["channel_format"], kwargs["num_channels"])]
-                kwargs["fmt"] = fmt
-                raise TaichiRuntimeError(
-                    "The channel_format and num_channels arguments for texture are deprecated in v1.6.0, "
-                    "and they are removed in v1.7.0. Please use fmt instead."
-                )
-            else:
-                raise TaichiRuntimeError(
-                    "The channel_format and num_channels arguments are no longer required for non-RW textures "
-                    "since v1.6.0, and they are removed in v1.7.0. Please remove them."
-                )
-
-
-def _check_args(kwargs: Dict[str, Any], allowed_kwargs: List[str]):
-    for k, v in kwargs.items():
-        if k not in allowed_kwargs:
-            raise TaichiRuntimeError(
-                f"Invalid argument: {k}, you can only create a graph argument with: {allowed_kwargs}"
-            )
-        if k == "tag":
-            if not isinstance(v, ArgKind):
-                raise TaichiRuntimeError(f"tag must be a ArgKind variant, but found {type(v)}.")
-        if k == "name":
-            if not isinstance(v, str):
-                raise TaichiRuntimeError(f"name must be a string, but found {type(v)}.")
-
-
-def _make_arg_scalar(kwargs: Dict[str, Any]):
-    allowed_kwargs = [
-        "tag",
-        "name",
-        "dtype",
-    ]
-    _check_args(kwargs, allowed_kwargs)
-    name = kwargs["name"]
-    dtype = kwargs["dtype"]
-    if isinstance(dtype, MatrixType):
-        raise TaichiRuntimeError(f"Tag ArgKind.SCALAR must specify a scalar type, but found {type(dtype)}.")
-    return _ti_core.Arg(ArgKind.SCALAR, name, dtype, 0, [])
-
-
-def _make_arg_ndarray(kwargs: Dict[str, Any]):
-    allowed_kwargs = [
-        "tag",
-        "name",
-        "dtype",
-        "ndim",
-        "element_shape",
-    ]
-    _check_args(kwargs, allowed_kwargs)
-    name = kwargs["name"]
-    ndim = kwargs["ndim"]
-    dtype = kwargs["dtype"]
-    element_shape = kwargs["element_shape"]
-    if isinstance(dtype, MatrixType):
-        raise TaichiRuntimeError(f"Tag ArgKind.NDARRAY must specify a scalar type, but found {dtype}.")
-    return _ti_core.Arg(ArgKind.NDARRAY, name, dtype, ndim, element_shape)
-
-
-def _make_arg_matrix(kwargs: Dict[str, Any]):
-    allowed_kwargs = [
-        "tag",
-        "name",
-        "dtype",
-    ]
-    _check_args(kwargs, allowed_kwargs)
-    name = kwargs["name"]
-    dtype = kwargs["dtype"]
-    if not isinstance(dtype, MatrixType):
-        raise TaichiRuntimeError(f"Tag ArgKind.MATRIX must specify matrix type, but got {dtype}.")
-    return _ti_core.Arg(ArgKind.MATRIX, f"{name}", dtype.dtype, 0, [dtype.n, dtype.m])
-
-
-def _make_arg_texture(kwargs: Dict[str, Any]):
-    allowed_kwargs = [
-        "tag",
-        "name",
-        "ndim",
-    ]
-    _check_args(kwargs, allowed_kwargs)
-    name = kwargs["name"]
-    ndim = kwargs["ndim"]
-    return _ti_core.Arg(ArgKind.TEXTURE, name, impl.f32, 4, [2] * ndim)
-
-
-def _make_arg_rwtexture(kwargs: Dict[str, Any]):
-    allowed_kwargs = [
-        "tag",
-        "name",
-        "ndim",
-        "fmt",
-    ]
-    _check_args(kwargs, allowed_kwargs)
-    name = kwargs["name"]
-    ndim = kwargs["ndim"]
-    fmt = kwargs["fmt"]
-    if fmt == enums.Format.unknown:
-        raise TaichiRuntimeError(f"Tag ArgKind.RWTEXTURE must specify a valid color format, but found {fmt}.")
-    channel_format, num_channels = FORMAT2TY_CH[fmt]
-    return _ti_core.Arg(ArgKind.RWTEXTURE, name, channel_format, num_channels, [2] * ndim)
-
-
-def _make_arg(kwargs: Dict[str, Any]):
-    assert "tag" in kwargs
-    _deprecate_arg_args(kwargs)
-    proc = {
-        ArgKind.SCALAR: _make_arg_scalar,
-        ArgKind.NDARRAY: _make_arg_ndarray,
-        ArgKind.MATRIX: _make_arg_matrix,
-        ArgKind.TEXTURE: _make_arg_texture,
-        ArgKind.RWTEXTURE: _make_arg_rwtexture,
-    }
-    tag = kwargs["tag"]
-    return proc[tag](kwargs)
-
-
-def _kwarg_rewriter(args, kwargs):
-    for i, arg in enumerate(args):
-        rewrite_map = {
-            0: "tag",
-            1: "name",
-            2: "dtype",
-            3: "ndim",
-            4: "field_dim",
-            5: "element_shape",
-            6: "channel_format",
-            7: "shape",
-            8: "num_channels",
-        }
-        if i in rewrite_map:
-            kwargs[rewrite_map[i]] = arg
-        else:
-            raise TaichiRuntimeError(f"Unexpected {i}th positional argument")
-
-
-def Arg(*args, **kwargs):
-    _kwarg_rewriter(args, kwargs)
-    return _make_arg(kwargs)
-
-
-__all__ = ["GraphBuilder", "Graph", "Arg", "ArgKind"]