python/tvm_ffi/core.pyi

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Typestubs for Cython."""

from __future__ import annotations

import types
from ctypes import c_void_p
from enum import IntEnum
from typing import Any, Callable

# Public module-level variables referenced by Python code
MISSING: Object
KWARGS: Object
ERROR_NAME_TO_TYPE: dict[str, type]
ERROR_TYPE_TO_NAME: dict[type, str]

_WITH_APPEND_BACKTRACE: Callable[[BaseException, str], BaseException] | None
_TRACEBACK_TO_BACKTRACE_STR: Callable[[types.TracebackType | None], str] | None
# DLPack protocol version (defined in tensor.pxi)
__dlpack_version__: tuple[int, int]

class CObject:
    def __ctypes_handle__(self) -> Any: ...
    def __chandle__(self) -> int: ...
    @property
    def id_(self) -> int: ...
    def __reduce__(self) -> Any: ...
    def __getstate__(self) -> dict[str, Any]: ...
    def __setstate__(self, state: dict[str, Any]) -> None: ...
    def __repr__(self) -> str: ...
    def __eq__(self, other: Any) -> bool: ...
    def __ne__(self, other: Any) -> bool: ...
    def __hash__(self) -> int: ...
    def __init_handle_by_constructor__(self, fconstructor: Any, *args: Any) -> None: ...
    def same_as(self, other: Any) -> bool: ...
    def is_(self, other: Any) -> bool: ...
    def _move(self) -> ObjectRValueRef: ...
    def __move_handle_from__(self, other: CObject) -> None: ...

class CContainerBase(CObject): ...
class Object(CObject): ...

def object_repr(obj: CObject) -> str: ...

class ObjectConvertible:
    def asobject(self) -> Object: ...

class ObjectRValueRef:
    obj: Object
    def __init__(self, obj: Object) -> None: ...

class OpaquePyObject(Object):
    def pyobject(self) -> Any: ...

class PyNativeObject:
    __slots__: list[str]
    def __init_cached_object_by_constructor__(self, fconstructor: Any, *args: Any) -> None: ...

def _register_object_by_index(type_index: int, type_cls: type) -> TypeInfo: ...
def _object_type_key_to_index(type_key: str) -> int | None: ...
def _set_type_cls(type_info: TypeInfo, type_cls: type) -> None: ...
def _lookup_or_register_type_info_from_type_key(type_key: str) -> TypeInfo: ...
def _lookup_type_attr(type_index: int, attr_key: str) -> Any: ...
def _register_type_attr(type_index: int, attr_key: str, value: Any) -> None: ...
def _type_cls_to_type_info(type_cls: type) -> TypeInfo | None: ...

class Error(Object):
    def __init__(self, kind: str, message: str, backtrace: str) -> None: ...
    def update_backtrace(self, backtrace: str) -> None: ...
    def py_error(self) -> BaseException: ...
    @property
    def kind(self) -> str: ...
    @property
    def message(self) -> str: ...
    @property
    def backtrace(self) -> str: ...

def _convert_to_ffi_error(error: BaseException) -> Error: ...
def _env_set_current_stream(
    device_type: int, device_id: int, stream: int | c_void_p
) -> int | c_void_p: ...
def _env_get_current_stream(device_type: int, device_id: int) -> int: ...

class DataType:
    def __init__(self, dtype_str: str) -> None: ...
    def __reduce__(self) -> Any: ...
    def __eq__(self, other: Any) -> bool: ...
    def __ne__(self, other: Any) -> bool: ...
    @property
    def type_code(self) -> int: ...
    @property
    def bits(self) -> int: ...
    @property
    def lanes(self) -> int: ...
    @property
    def itemsize(self) -> int: ...
    def __str__(self) -> str: ...

def _set_class_dtype(cls: type) -> None: ...
def _convert_torch_dtype_to_ffi_dtype(torch_dtype: Any) -> DataType: ...
def _convert_numpy_dtype_to_ffi_dtype(numpy_dtype: Any) -> DataType: ...
def _create_cdtype_from_tuple(
    cls: type[DataType], code: int, bits: int, lanes: int
) -> DataType: ...

class DLDeviceType(IntEnum):
    kDLCPU = 1
    kDLCUDA = 2
    kDLCUDAHost = 3
    kDLOpenCL = 4
    kDLVulkan = 7
    kDLMetal = 8
    kDLVPI = 9
    kDLROCM = 10
    kDLROCMHost = 11
    kDLExtDev = 12
    kDLCUDAManaged = 13
    kDLOneAPI = 14
    kDLWebGPU = 15
    kDLHexagon = 16
    kDLTrn = 17

class Device:
    def __init__(self, device_type: str | int, index: int | None = None) -> None: ...
    def __reduce__(self) -> Any: ...
    def __eq__(self, other: Any) -> bool: ...
    def __ne__(self, other: Any) -> bool: ...
    def __str__(self) -> str: ...
    def __repr__(self) -> str: ...
    def __hash__(self) -> int: ...
    def __device_type_name__(self) -> str: ...
    @property
    def type(self) -> str: ...
    @property
    def index(self) -> int: ...
    def dlpack_device_type(self) -> int: ...

def _set_class_device(cls: type) -> None: ...

_CLASS_DEVICE: type[Device]

def _shape_obj_get_py_tuple(obj: Any) -> tuple[int, ...]: ...

class Tensor(Object):
    @property
    def shape(self) -> tuple[int, ...]: ...
    @property
    def strides(self) -> tuple[int, ...]: ...
    @property
    def dtype(self) -> Any: ...
    @property
    def device(self) -> Device: ...
    def _to_dlpack(self) -> Any: ...
    def _to_dlpack_versioned(self) -> Any: ...
    def __dlpack_device__(self) -> tuple[int, int]: ...
    def __dlpack__(
        self,
        *,
        stream: Any | None = None,
        max_version: tuple[int, int] | None = None,
        dl_device: tuple[int, int] | None = None,
        copy: bool | None = None,
    ) -> Any: ...

_CLASS_TENSOR: type[Tensor] = Tensor

def _set_class_tensor(cls: type[Tensor]) -> None: ...
def from_dlpack(
    ext_tensor: Any, *, require_alignment: int = ..., require_contiguous: bool = ...
) -> Tensor: ...

class DLTensorTestWrapper:
    __dlpack_c_exchange_api__: int
    def __init__(self, tensor: Tensor) -> None: ...
    def __tvm_ffi_env_stream__(self) -> int: ...
    def __dlpack_device__(self) -> tuple[int, int]: ...
    def __dlpack__(self, **kwargs: Any) -> Any: ...

def _dltensor_test_wrapper_c_dlpack_from_pyobject_as_intptr() -> int: ...

class Function(Object):
    @property
    def release_gil(self) -> bool: ...
    @release_gil.setter
    def release_gil(self, value: bool) -> None: ...
    def __call__(self, *args: Any) -> Any: ...
    @staticmethod
    def __from_extern_c__(c_symbol: int, *, keep_alive_object: Any | None = None) -> Function: ...
    @staticmethod
    def __from_mlir_packed_safe_call__(
        mlir_packed_symbol: int, *, keep_alive_object: Any | None = None
    ) -> Function: ...

def _register_global_func(
    name: str, pyfunc: Callable[..., Any] | Function, override: bool
) -> Function: ...
def _get_global_func(name: str, allow_missing: bool) -> Function | None: ...
def _convert_to_ffi_func(pyfunc: Callable[..., Any], tensor_cls: type | None = ...) -> Function: ...
def _convert_to_opaque_object(pyobject: Any) -> OpaquePyObject: ...
def _print_debug_info() -> None: ...
def _register_py_class(parent_type_info: TypeInfo, type_key: str, type_cls: type) -> TypeInfo: ...
def _register_fields(type_info: TypeInfo, fields: list[Any]) -> list[TypeField]: ...

class String(str, PyNativeObject):
    __slots__ = ["_tvm_ffi_cached_object"]
    _tvm_ffi_cached_object: Object | None

    def __new__(cls, value: str) -> String: ...

    # pylint: disable=no-self-argument
    def __from_tvm_ffi_object__(cls, obj: Any) -> String: ...

class Bytes(bytes, PyNativeObject):
    __slots__ = ["_tvm_ffi_cached_object"]
    _tvm_ffi_cached_object: Object | None

    def __new__(cls, value: bytes) -> Bytes: ...

    # pylint: disable=no-self-argument
    def __from_tvm_ffi_object__(cls, obj: Any) -> Bytes: ...

# ---------------------------------------------------------------------------
# Owned FFI value container (from cython/object.pxi)
# ---------------------------------------------------------------------------

class CAny:
    def __init__(self, value: Any = None) -> None: ...
    @property
    def type_index(self) -> int: ...

def _to_py_class_value(value: CAny) -> Any: ...

# ---------------------------------------------------------------------------
# Type reflection metadata (from cython/type_info.pxi)
# ---------------------------------------------------------------------------

class TypeSchema:
    origin: str
    args: tuple[TypeSchema, ...] = ()
    origin_type_index: int

    def __init__(
        self,
        origin: str,
        args: tuple[TypeSchema, ...] = (),
        origin_type_index: int = ...,
    ) -> None: ...
    @staticmethod
    def from_json_obj(obj: dict[str, Any]) -> TypeSchema: ...
    @staticmethod
    def from_json_str(s: str) -> TypeSchema: ...
    @staticmethod
    def from_type_index(type_index: int, args: tuple[TypeSchema, ...] = ()) -> TypeSchema: ...
    @staticmethod
    def from_annotation(annotation: object) -> TypeSchema: ...
    def repr(self, ty_map: Callable[[str], str] | None = None) -> str: ...
    def check_value(self, value: object) -> None: ...
    def convert(self, value: object) -> CAny: ...
    def to_json(self) -> dict[str, Any]: ...

class TypeField:
    name: str
    doc: str | None
    size: int
    offset: int
    frozen: bool
    metadata: dict[str, Any]
    getter: Any
    setter: Any
    ty: TypeSchema | None
    c_init: bool
    c_kw_only: bool
    c_has_default: bool
    c_default: Any
    c_default_factory: Any
    c_repr: bool
    c_compare: bool
    c_hash: bool
    c_structural_eq: str | None
    dataclass_field: Any | None

    def as_property(self, cls: type) -> property: ...

class TypeMethod:
    name: str
    doc: str | None
    func: Any
    is_static: bool
    metadata: dict[str, Any]

    def as_callable(self, cls: type) -> Callable[..., Any]: ...

class TypeInfo:
    type_cls: type | None
    type_index: int
    type_key: str
    type_ancestors: list[int]
    fields: list[TypeField]
    methods: list[TypeMethod]
    parent_type_info: TypeInfo | None

    def _register_fields(self, fields: list[Any]) -> None: ...
    def prototype_py(self) -> str: ...