gstaichi 0.1.25.dev0__cp313-cp313-win_amd64.whl

gstaichi/lang/_template_mapper.py

@@ -0,0 +1,199 @@
+import dataclasses
+import weakref
+from typing import Any, Union
+
+import gstaichi.lang
+import gstaichi.lang._ndarray
+import gstaichi.lang._texture
+import gstaichi.lang.expr
+import gstaichi.lang.snode
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang.any_array import AnyArray
+from gstaichi.lang.argpack import ArgPack, ArgPackType
+from gstaichi.lang.exception import (
+    GsTaichiRuntimeTypeError,
+)
+from gstaichi.lang.kernel_arguments import KernelArgument
+from gstaichi.lang.matrix import MatrixType
+from gstaichi.lang.util import to_gstaichi_type
+from gstaichi.types import (
+    ndarray_type,
+    sparse_matrix_builder,
+    template,
+    texture_type,
+)
+
+AnnotationType = Union[
+    template,
+    ArgPackType,
+    "texture_type.TextureType",
+    "texture_type.RWTextureType",
+    ndarray_type.NdarrayType,
+    sparse_matrix_builder,
+    Any,
+]
+
+
+class GsTaichiCallableTemplateMapper:
+    """
+    This should probably be renamed to sometihng like FeatureMapper, or
+    FeatureExtractor, since:
+    - it's not specific to templates
+    - it extracts what are later called 'features', for example for ndarray this includes:
+        - element type
+        - number dimensions
+        - needs grad (or not)
+    - these are returned as a heterogeneous tuple, whose contents depends on the type
+    """
+
+    def __init__(self, arguments: list[KernelArgument], template_slot_locations: list[int]) -> None:
+        self.arguments: list[KernelArgument] = arguments
+        self.num_args: int = len(arguments)
+        self.template_slot_locations: list[int] = template_slot_locations
+        self.mapping: dict[tuple[Any, ...], int] = {}
+
+    @staticmethod
+    def extract_arg(arg, annotation: AnnotationType, arg_name: str) -> Any:
+        if annotation == template or isinstance(annotation, template):
+            if isinstance(arg, gstaichi.lang.snode.SNode):
+                return arg.ptr
+            if isinstance(arg, gstaichi.lang.expr.Expr):
+                return arg.ptr.get_underlying_ptr_address()
+            if isinstance(arg, _ti_core.ExprCxx):
+                return arg.get_underlying_ptr_address()
+            if isinstance(arg, tuple):
+                return tuple(GsTaichiCallableTemplateMapper.extract_arg(item, annotation, arg_name) for item in arg)
+            if isinstance(arg, gstaichi.lang._ndarray.Ndarray):
+                raise GsTaichiRuntimeTypeError(
+                    "Ndarray shouldn't be passed in via `ti.template()`, please annotate your kernel using `ti.types.ndarray(...)` instead"
+                )
+
+            if isinstance(arg, (list, tuple, dict, set)) or hasattr(arg, "_data_oriented"):
+                # [Composite arguments] Return weak reference to the object
+                # GsTaichi kernel will cache the extracted arguments, thus we can't simply return the original argument.
+                # Instead, a weak reference to the original value is returned to avoid memory leak.
+
+                # TODO(zhanlue): replacing "tuple(args)" with "hash of argument values"
+                # This can resolve the following issues:
+                # 1. Invalid weak-ref will leave a dead(dangling) entry in both caches: "self.mapping" and "self.compiled_functions"
+                # 2. Different argument instances with same type and same value, will get templatized into seperate kernels.
+                return weakref.ref(arg)
+
+            # [Primitive arguments] Return the value
+            return arg
+        if isinstance(annotation, ArgPackType):
+            if not isinstance(arg, ArgPack):
+                raise GsTaichiRuntimeTypeError(f"Argument {arg_name} must be a argument pack, got {type(arg)}")
+            return tuple(
+                GsTaichiCallableTemplateMapper.extract_arg(arg[name], dtype, arg_name)
+                for index, (name, dtype) in enumerate(annotation.members.items())
+            )
+        if dataclasses.is_dataclass(annotation):
+            _res_l = []
+            for field in dataclasses.fields(annotation):
+                field_value = getattr(arg, field.name)
+                arg_name = f"__ti_{arg_name}_{field.name}"
+                field_extracted = GsTaichiCallableTemplateMapper.extract_arg(field_value, field.type, arg_name)
+                _res_l.append(field_extracted)
+            return tuple(_res_l)
+        if isinstance(annotation, texture_type.TextureType):
+            if not isinstance(arg, gstaichi.lang._texture.Texture):
+                raise GsTaichiRuntimeTypeError(f"Argument {arg_name} must be a texture, got {type(arg)}")
+            if arg.num_dims != annotation.num_dimensions:
+                raise GsTaichiRuntimeTypeError(
+                    f"TextureType dimension mismatch for argument {arg_name}: expected {annotation.num_dimensions}, got {arg.num_dims}"
+                )
+            return (arg.num_dims,)
+        if isinstance(annotation, texture_type.RWTextureType):
+            if not isinstance(arg, gstaichi.lang._texture.Texture):
+                raise GsTaichiRuntimeTypeError(f"Argument {arg_name} must be a texture, got {type(arg)}")
+            if arg.num_dims != annotation.num_dimensions:
+                raise GsTaichiRuntimeTypeError(
+                    f"RWTextureType dimension mismatch for argument {arg_name}: expected {annotation.num_dimensions}, got {arg.num_dims}"
+                )
+            if arg.fmt != annotation.fmt:
+                raise GsTaichiRuntimeTypeError(
+                    f"RWTextureType format mismatch for argument {arg_name}: expected {annotation.fmt}, got {arg.fmt}"
+                )
+            # (penguinliong) '0' is the assumed LOD level. We currently don't
+            # support mip-mapping.
+            return arg.num_dims, arg.fmt, 0
+        if isinstance(annotation, ndarray_type.NdarrayType):
+            if isinstance(arg, gstaichi.lang._ndarray.Ndarray):
+                annotation.check_matched(arg.get_type(), arg_name)
+                needs_grad = (arg.grad is not None) if annotation.needs_grad is None else annotation.needs_grad
+                assert arg.shape is not None
+                return arg.element_type, len(arg.shape), needs_grad, annotation.boundary
+            if isinstance(arg, AnyArray):
+                ty = arg.get_type()
+                annotation.check_matched(arg.get_type(), arg_name)
+                return ty.element_type, len(arg.shape), ty.needs_grad, annotation.boundary
+            # external arrays
+            shape = getattr(arg, "shape", None)
+            if shape is None:
+                raise GsTaichiRuntimeTypeError(f"Invalid type for argument {arg_name}, got {arg}")
+            shape = tuple(shape)
+            element_shape: tuple[int, ...] = ()
+            dtype = to_gstaichi_type(arg.dtype)
+            if isinstance(annotation.dtype, MatrixType):
+                if annotation.ndim is not None:
+                    if len(shape) != annotation.dtype.ndim + annotation.ndim:
+                        raise ValueError(
+                            f"Invalid value for argument {arg_name} - required array has ndim={annotation.ndim} element_dim={annotation.dtype.ndim}, "
+                            f"array with {len(shape)} dimensions is provided"
+                        )
+                else:
+                    if len(shape) < annotation.dtype.ndim:
+                        raise ValueError(
+                            f"Invalid value for argument {arg_name} - required element_dim={annotation.dtype.ndim}, "
+                            f"array with {len(shape)} dimensions is provided"
+                        )
+                element_shape = shape[-annotation.dtype.ndim :]
+                anno_element_shape = annotation.dtype.get_shape()
+                if None not in anno_element_shape and element_shape != anno_element_shape:
+                    raise ValueError(
+                        f"Invalid value for argument {arg_name} - required element_shape={anno_element_shape}, "
+                        f"array with element shape of {element_shape} is provided"
+                    )
+            elif annotation.dtype is not None:
+                # User specified scalar dtype
+                if annotation.dtype != dtype:
+                    raise ValueError(
+                        f"Invalid value for argument {arg_name} - required array has dtype={annotation.dtype.to_string()}, "
+                        f"array with dtype={dtype.to_string()} is provided"
+                    )
+
+                if annotation.ndim is not None and len(shape) != annotation.ndim:
+                    raise ValueError(
+                        f"Invalid value for argument {arg_name} - required array has ndim={annotation.ndim}, "
+                        f"array with {len(shape)} dimensions is provided"
+                    )
+            needs_grad = (
+                getattr(arg, "requires_grad", False) if annotation.needs_grad is None else annotation.needs_grad
+            )
+            element_type = (
+                _ti_core.get_type_factory_instance().get_tensor_type(element_shape, dtype)
+                if len(element_shape) != 0
+                else arg.dtype
+            )
+            return element_type, len(shape) - len(element_shape), needs_grad, annotation.boundary
+        if isinstance(annotation, sparse_matrix_builder):
+            return arg.dtype
+        # Use '#' as a placeholder because other kinds of arguments are not involved in template instantiation
+        return "#"
+
+    def extract(self, args: tuple[Any, ...]) -> tuple[Any, ...]:
+        extracted: list[Any] = []
+        for arg, kernel_arg in zip(args, self.arguments):
+            extracted.append(self.extract_arg(arg, kernel_arg.annotation, kernel_arg.name))
+        return tuple(extracted)
+
+    def lookup(self, args: tuple[Any, ...]) -> tuple[int, tuple[Any, ...]]:
+        if len(args) != self.num_args:
+            raise TypeError(f"{self.num_args} argument(s) needed but {len(args)} provided.")
+
+        key = self.extract(args)
+        if key not in self.mapping:
+            count = len(self.mapping)
+            self.mapping[key] = count
+        return self.mapping[key], key

gstaichi/lang/_texture.py

@@ -0,0 +1,172 @@
+# type: ignore
+
+import numpy as np
+
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang import impl
+from gstaichi.lang.expr import Expr, make_expr_group
+from gstaichi.lang.matrix import Matrix
+from gstaichi.lang.util import gstaichi_scope
+from gstaichi.types import vector
+from gstaichi.types.primitive_types import f32
+
+
+def _get_entries(mat):
+    if isinstance(mat, Matrix):
+        return mat.entries
+    return [mat]
+
+
+class TextureSampler:
+    def __init__(self, ptr_expr, num_dims) -> None:
+        self.ptr_expr = ptr_expr
+        self.num_dims = num_dims
+
+    @gstaichi_scope
+    def sample_lod(self, uv, lod):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(uv), lod)
+        v = ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kSampleLod, self.ptr_expr, args_group, dbg_info)
+        r = impl.call_internal("composite_extract_0", v, with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1", v, with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2", v, with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3", v, with_runtime_context=False)
+        return vector(4, f32)([r, g, b, a])
+
+    @gstaichi_scope
+    def fetch(self, index, lod):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(index), lod)
+        v = ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kFetchTexel, self.ptr_expr, args_group, dbg_info)
+        r = impl.call_internal("composite_extract_0", v, with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1", v, with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2", v, with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3", v, with_runtime_context=False)
+        return vector(4, f32)([r, g, b, a])
+
+
+class RWTextureAccessor:
+    def __init__(self, ptr_expr, num_dims) -> None:
+        # gstaichi_python.TexturePtrExpression.
+        self.ptr_expr = ptr_expr
+        self.num_dims = num_dims
+
+    @gstaichi_scope
+    def load(self, index):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(index))
+        v = ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kLoad, self.ptr_expr, args_group, dbg_info)
+        r = impl.call_internal("composite_extract_0", v, with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1", v, with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2", v, with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3", v, with_runtime_context=False)
+        return vector(4, f32)([r, g, b, a])
+
+    @gstaichi_scope
+    def store(self, index, value):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        args_group = make_expr_group(*_get_entries(index), *_get_entries(value))
+        impl.expr_init(
+            ast_builder.make_texture_op_expr(_ti_core.TextureOpType.kStore, self.ptr_expr, args_group, dbg_info)
+        )
+
+    @property
+    @gstaichi_scope
+    def shape(self):
+        """A list containing sizes for each dimension. Note that element shape will be excluded.
+
+        Returns:
+            List[Int]: The result list.
+        """
+        dim = _ti_core.get_external_tensor_dim(self.ptr_expr)
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        ret = [Expr(_ti_core.get_external_tensor_shape_along_axis(self.ptr_expr, i, dbg_info)) for i in range(dim)]
+        return ret
+
+    @gstaichi_scope
+    def _loop_range(self):
+        """Gets the corresponding gstaichi_python.Expr to serve as loop range.
+
+        Returns:
+            gstaichi_python.Expr: See above.
+        """
+        return self.ptr_expr
+
+
+class Texture:
+    """GsTaichi Texture class.
+
+    Args:
+        fmt (ti.Format): Color format of the texture.
+        shape (Tuple[int]): Shape of the Texture.
+    """
+
+    def __init__(self, fmt, arr_shape):
+        self.tex = impl.get_runtime().prog.create_texture(fmt, arr_shape)
+        self.fmt = fmt
+        self.num_dims = len(arr_shape)
+        self.shape = arr_shape
+
+    def from_ndarray(self, ndarray):
+        """Loads an ndarray to texture.
+
+        Args:
+            ndarray (ti.Ndarray): Source ndarray to load from.
+        """
+        self.tex.from_ndarray(ndarray.arr)
+
+    def from_field(self, field):
+        """Loads a field to texture.
+
+        Args:
+            field (ti.Field): Source field to load from.
+        """
+        self.tex.from_snode(field.snode.ptr)
+
+    def _device_allocation_ptr(self):
+        return self.tex.device_allocation_ptr()
+
+    def from_image(self, image):
+        """Loads a PIL image to texture. This method is only allowed a 2D texture with `ti.Format.rgba8`.
+
+        Args:
+            image (PIL.Image.Image): Source PIL image to load from.
+
+        """
+        from PIL import Image  # pylint: disable=import-outside-toplevel
+
+        assert isinstance(image, Image.Image)
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        assert image.size == tuple(self.shape)
+
+        assert self.num_dims == 2
+        # Don't use transpose method since its enums are too new
+        image = image.rotate(90, expand=True)
+        arr = np.asarray(image)
+        from gstaichi._kernels import (  # pylint: disable=import-outside-toplevel
+            load_texture_from_numpy,
+        )
+
+        load_texture_from_numpy(self, arr)
+
+    def to_image(self):
+        """Saves a texture to a PIL image in RGB mode. This method is only allowed a 2D texture with `ti.Format.rgba8`.
+
+        Returns:
+            img (PIL.Image.Image): a PIL image in RGB mode, with the same size as source texture.
+        """
+        assert self.num_dims == 2
+        from PIL import Image  # pylint: disable=import-outside-toplevel
+
+        res = np.zeros(self.shape + (3,), np.uint8)
+        from gstaichi._kernels import (  # pylint: disable=import-outside-toplevel
+            save_texture_to_numpy,
+        )
+
+        save_texture_to_numpy(self, res)
+        return Image.fromarray(res).rotate(270, expand=True)

gstaichi/lang/_wrap_inspect.py

@@ -0,0 +1,189 @@
+# type: ignore
+
+# GsTaichi's custom inspect module.
+# This module is used by GsTaichi's ast transformer to parse the source code.
+# Currently this module is aimed for working in the following modes:
+# 1. Usual Python/IPython mode, e.g. python script.py
+#    In this case we mainly rely on the built-in `inspect` module, except
+#    we need some hacks when we are in IPython mode and there is a cell magic.
+# 2. Blender's scripting mode, e.g. Users write GsTaichi code in the scripting
+#    window in Blender and press the run button. In this case we need to
+#    retrieve the source using Blender's `bpy.data.texts` and write it to a temp
+#    file so that the inspect module can parse.
+# 3. The interactive shell mode, e.g. Users directly type their code in the
+#    interactive shell. In this case we use `dill` to get the source.
+#
+# NB: Running GsTaichi in other modes are likely not supported.
+
+import atexit
+import inspect
+import os
+import tempfile
+
+import dill
+
+_builtin_getfile = inspect.getfile
+_builtin_findsource = inspect.findsource
+
+
+def _find_source_with_custom_getfile_func(func, obj):
+    """Use a custom function `func` to replace inspect's `getfile`, return the
+    source found by the new routine and restore the original `getfile` back.
+    """
+    inspect.getfile = func  # replace with our custom func
+    source = inspect.findsource(obj)
+    inspect.getfile = _builtin_getfile  # restore
+    return source
+
+
+def _blender_get_text_name(filename: str):
+    """Extract filename from path in the Blender mode."""
+    # In Blender's scripting mode, unsaved files are named
+    # like `/Text`, `/Text.001`, `/test.py`, etc.
+    # We simply remove this path seperator.
+    if filename.startswith(os.path.sep) and filename.count(os.path.sep) == 1:
+        return filename[1:]  # "/Text.001" --> "Text.001"
+
+    # Saved text files are named like `some-path/xxx.blend/Text` or
+    # `some-path/xxx.blend/test.py`
+    # We drop the path and extract the filename with extension.
+    index = filename.rfind(".blend" + os.path.sep)
+    if index != -1:
+        return filename[index + 7 :]  # "xxx.blend/test.py" --> "test.py"
+
+    return None
+
+
+def _blender_findsource(obj):
+    try:
+        import bpy  # pylint: disable=import-outside-toplevel
+    except:
+        raise ImportError("Not in Blender environment!")
+
+    # Inspect's built-in `getfile` returns the filename like
+    # `/Text`, `/Text.001`, `some-path/xxx.blend/test.py`
+    # This filename may not be a full valid path.
+    filename = _builtin_getfile(obj)
+    # Extract the text name without path
+    text_name = _blender_get_text_name(filename)
+    if text_name is None:
+        raise IOError("Object `{obj.__name__}` is not defined in a .blend file!")
+    # Get the lines of code via text_name
+    lines = bpy.data.texts[text_name].as_string()
+    # Now we have found the lines of code.
+    # We first check if they are already cached, to avoid file io in each query.
+    try:
+        filename = _blender_findsource._saved_inspect_cache[lines]  # pylint: disable=no-member
+    except KeyError:
+        # Save the code to a valid path.
+        fd, filename = tempfile.mkstemp(prefix="_Blender_", suffix=f"_{text_name}.py")
+        os.close(fd)
+
+        with open(filename, "w") as f:
+            f.write(lines)
+
+        _blender_findsource._saved_inspect_cache[lines] = filename  # pylint: disable=no-member
+        atexit.register(os.unlink, filename)  # Remove file when program exits
+
+    # Our custom getfile function
+    def wrapped_getfile(ob):
+        if id(ob) == id(obj):
+            return filename
+
+        return _builtin_getfile(ob)
+
+    return _find_source_with_custom_getfile_func(wrapped_getfile, obj)
+
+
+_blender_findsource._saved_inspect_cache = {}
+
+
+def _Python_IPython_findsource(obj):
+    try:
+        # In Python and IPython the builtin inspect would suffice in most cases
+        return _builtin_findsource(obj)
+    except IOError:
+        # Except that the cell has a magic command like %%time or %%timeit
+        # In this case the filename returned by the built-in's getfile is wrong,
+        # it becomes something like `<timed exec>` or `<magic-timeit>`.
+        filename = _builtin_getfile(obj)
+        if filename in {"<timed exec>", "<magic-timeit>"}:
+            try:
+                ip = get_ipython()
+                if ip is not None:
+                    # So we are in IPython's cell magic
+                    session_id = ip.history_manager.get_last_session_id()
+                    fd, filename = tempfile.mkstemp(prefix="_IPython_", suffix=f"_{session_id}.py")
+                    os.close(fd)
+                    # The latest lines of code can be retrived from here
+                    lines = ip.history_manager._i00
+
+                    # `lines` is a string that also contains the cell magic
+                    # command, we need to remove the magic command
+                    # (and spaces/sep around it) to obtain a valid Python code
+                    # snippet before saving it to a file
+                    index = lines.find("%time")
+                    lines_stripped = lines[index:]
+                    lines_stripped = lines_stripped.split(maxsplit=1)[1]
+
+                    with open(filename, "w") as f:
+                        f.write(lines_stripped)
+
+                    atexit.register(os.unlink, filename)  # Remove the file after the program exits
+                    func = lambda obj: filename
+                    return _find_source_with_custom_getfile_func(func, obj)
+
+            except ImportError:
+                pass
+        raise IOError(
+            f"Cannot find source code for Object: {obj}, it's likely \
+you are not running GsTaichi from command line or IPython."
+        )
+
+
+def _REPL_findsource(obj):
+    """Findsource in the interactive shell mode."""
+    return dill.source.findsource(obj)
+
+
+def _custom_findsource(obj):
+    try:
+        return _Python_IPython_findsource(obj)
+    except IOError:
+        try:
+            return _REPL_findsource(obj)
+        except:
+            try:
+                return _blender_findsource(obj)
+            except:
+                raise IOError(
+                    f"Cannot find source code for Object: {obj}, this \
+is possibly because of you are running GsTaichi in an environment that GsTaichi's own \
+inspect module cannot find the source. Please report an issue to help us fix: \
+https://github.com/taichi-dev/gstaichi/issues"
+                )
+
+
+class _InspectContextManager:
+    def __enter__(self):
+        inspect.findsource = _custom_findsource
+        return self
+
+    def __exit__(self, *_):
+        inspect.findsource = _builtin_findsource
+
+
+def getsourcelines(obj):
+    with _InspectContextManager():
+        return inspect.getsourcelines(obj)
+
+
+def getsourcefile(obj):
+    with _InspectContextManager():
+        ret = inspect.getsourcefile(obj)
+        if ret is None:
+            ret = inspect.getfile(obj)
+        return ret
+
+
+__all__ = ["getsourcelines", "getsourcefile"]

gstaichi/lang/any_array.py

@@ -0,0 +1,99 @@
+# type: ignore
+
+from gstaichi._lib import core as _ti_core
+from gstaichi.lang import impl
+from gstaichi.lang.expr import Expr, make_expr_group
+from gstaichi.lang.util import gstaichi_scope
+from gstaichi.types.enums import Layout
+from gstaichi.types.ndarray_type import NdarrayTypeMetadata
+
+
+class AnyArray:
+    """Class for arbitrary arrays in Python AST.
+
+    Args:
+        ptr (gstaichi_python.Expr): A gstaichi_python.Expr wrapping a gstaichi_python.ExternalTensorExpression.
+        element_shape (Tuple[Int]): () if scalar elements (default), (n) if vector elements, and (n, m) if matrix elements.
+        layout (Layout): Memory layout.
+    """
+
+    def __init__(self, ptr):
+        assert ptr.is_external_tensor_expr()
+        self.ptr = ptr
+        self.ptr.type_check(impl.get_runtime().prog.config())
+
+    def element_shape(self):
+        return _ti_core.get_external_tensor_element_shape(self.ptr)
+
+    def layout(self):
+        # 0: scalar; 1: vector (SOA); 2: matrix (SOA); -1: vector
+        # (AOS); -2: matrix (AOS)
+        element_dim = _ti_core.get_external_tensor_element_dim(self.ptr)
+        if element_dim == 1 or element_dim == 2:
+            return Layout.SOA
+        return Layout.AOS
+
+    def get_type(self):
+        return NdarrayTypeMetadata(
+            _ti_core.get_external_tensor_element_type(self.ptr), None, _ti_core.get_external_tensor_needs_grad(self.ptr)
+        )  # AnyArray can take any shape
+
+    @property
+    @gstaichi_scope
+    def grad(self):
+        """Returns the gradient of this array."""
+        return AnyArray(_ti_core.make_external_tensor_grad_expr(self.ptr))
+
+    @property
+    @gstaichi_scope
+    def shape(self):
+        """A list containing sizes for each dimension. Note that element shape will be excluded.
+
+        Returns:
+            List[Int]: The result list.
+        """
+        dim = _ti_core.get_external_tensor_dim(self.ptr)
+        dbg_info = _ti_core.DebugInfo(impl.get_runtime().get_current_src_info())
+        return [Expr(_ti_core.get_external_tensor_shape_along_axis(self.ptr, i, dbg_info)) for i in range(dim)]
+
+    @gstaichi_scope
+    def _loop_range(self):
+        """Gets the corresponding gstaichi_python.Expr to serve as loop range.
+
+        Returns:
+            gstaichi_python.Expr: See above.
+        """
+        return self.ptr
+
+
+class AnyArrayAccess:
+    """Class for first-level access to AnyArray with Vector/Matrix elements in Python AST.
+
+    Args:
+        arr (AnyArray): See above.
+        indices_first (Tuple[Int]): Indices of first-level access.
+    """
+
+    def __init__(self, arr, indices_first):
+        self.arr = arr
+        self.indices_first = indices_first
+
+    @gstaichi_scope
+    def subscript(self, i, j):
+        ast_builder = impl.get_runtime().compiling_callable.ast_builder()
+
+        indices_second = (i,) if len(self.arr.element_shape()) == 1 else (i, j)
+        if self.arr.layout() == Layout.SOA:
+            indices = indices_second + self.indices_first
+        else:
+            indices = self.indices_first + indices_second
+        return Expr(
+            ast_builder.expr_subscript(
+                self.arr.ptr,
+                make_expr_group(*indices),
+                _ti_core.DebugInfo(impl.get_runtime().get_current_src_info()),
+            )
+        )
+
+
+__all__ = []