onnx--onnx/onnx/onnx_cpp2py_export/defs.pyi

"""Submodule containing all the ONNX schema definitions."""

from collections.abc import Callable, Sequence
from typing import overload

from onnx import AttributeProto, FunctionProto
from onnx.shape_inference import InferenceContext

class SchemaError(Exception): ...

class OpSchema:
    def __init__(
        self,
        name: str,
        domain: str,
        since_version: int,
        doc: str = "",
        *,
        inputs: Sequence[OpSchema.FormalParameter] = (),
        outputs: Sequence[OpSchema.FormalParameter] = (),
        type_constraints: Sequence[tuple[str, Sequence[str], str]] = (),
        attributes: Sequence[OpSchema.Attribute] = (),
    ) -> None: ...
    @property
    def file(self) -> str: ...
    @property
    def line(self) -> int: ...
    @property
    def support_level(self) -> SupportType: ...
    @property
    def doc(self) -> str | None: ...
    @property
    def since_version(self) -> int: ...
    @property
    def deprecated(self) -> bool: ...
    @property
    def domain(self) -> str: ...
    @property
    def name(self) -> str: ...
    @property
    def min_input(self) -> int: ...
    @property
    def max_input(self) -> int: ...
    @property
    def min_output(self) -> int: ...
    @property
    def max_output(self) -> int: ...
    @property
    def attributes(self) -> dict[str, Attribute]: ...
    @property
    def inputs(self) -> Sequence[FormalParameter]: ...
    @property
    def outputs(self) -> Sequence[FormalParameter]: ...
    @property
    def type_constraints(self) -> Sequence[TypeConstraintParam]: ...
    @property
    def has_type_and_shape_inference_function(self) -> bool: ...
    @property
    def has_data_propagation_function(self) -> bool: ...
    @staticmethod
    def is_infinite(v: int) -> bool: ...
    def consumed(self, schema: OpSchema, i: int) -> tuple[UseType, int]: ...
    def _infer_node_outputs(
        self,
        node_proto: bytes,
        value_types: dict[str, bytes],
        input_data: dict[str, bytes],
        input_sparse_data: dict[str, bytes],
    ) -> dict[str, bytes]: ...
    @property
    def function_body(self) -> FunctionProto: ...
    def set_type_and_shape_inference_function(
        self, func: Callable[[InferenceContext], None]
    ) -> None: ...
    def get_type_and_shape_inference_function(
        self,
    ) -> Callable[[InferenceContext], None]: ...

    class TypeConstraintParam:
        def __init__(
            self,
            type_param_str: str,
            allowed_type_strs: Sequence[str],
            description: str = "",
        ) -> None:
            """Type constraint parameter.

            Args:
                type_param_str: Type parameter string, for example, "T", "T1", etc.
                allowed_type_strs: Allowed type strings for this type parameter. E.g. ["tensor(float)"].
                description: Type parameter description.
            """
        @property
        def type_param_str(self) -> str: ...
        @property
        def description(self) -> str: ...
        @property
        def allowed_type_strs(self) -> Sequence[str]: ...

    class FormalParameterOption:
        Single: OpSchema.FormalParameterOption = ...
        Optional: OpSchema.FormalParameterOption = ...
        Variadic: OpSchema.FormalParameterOption = ...

    class DifferentiationCategory:
        Unknown: OpSchema.DifferentiationCategory = ...
        Differentiable: OpSchema.DifferentiationCategory = ...
        NonDifferentiable: OpSchema.DifferentiationCategory = ...

    class NodeDeterminism:
        Deterministic: OpSchema.NodeDeterminism = ...
        NonDeterministic: OpSchema.NodeDeterminism = ...
        Unknown: OpSchema.NodeDeterminism = ...

    class FormalParameter:
        def __init__(
            self,
            name: str,
            type_str: str,
            description: str = "",
            *,
            param_option: OpSchema.FormalParameterOption = OpSchema.FormalParameterOption.Single,  # noqa: F821
            is_homogeneous: bool = True,
            min_arity: int = 1,
            differentiation_category: OpSchema.DifferentiationCategory = OpSchema.DifferentiationCategory.Unknown,  # noqa: F821
        ) -> None: ...
        @property
        def name(self) -> str: ...
        @property
        def types(self) -> set[str]: ...
        @property
        def type_str(self) -> str: ...
        @property
        def description(self) -> str: ...
        @property
        def option(self) -> OpSchema.FormalParameterOption: ...
        @property
        def is_homogeneous(self) -> bool: ...
        @property
        def min_arity(self) -> int: ...
        @property
        def differentiation_category(self) -> OpSchema.DifferentiationCategory: ...

    class AttrType:
        FLOAT: OpSchema.AttrType = ...
        INT: OpSchema.AttrType = ...
        STRING: OpSchema.AttrType = ...
        TENSOR: OpSchema.AttrType = ...
        GRAPH: OpSchema.AttrType = ...
        SPARSE_TENSOR: OpSchema.AttrType = ...
        TYPE_PROTO: OpSchema.AttrType = ...
        FLOATS: OpSchema.AttrType = ...
        INTS: OpSchema.AttrType = ...
        STRINGS: OpSchema.AttrType = ...
        TENSORS: OpSchema.AttrType = ...
        GRAPHS: OpSchema.AttrType = ...
        SPARSE_TENSORS: OpSchema.AttrType = ...
        TYPE_PROTOS: OpSchema.AttrType = ...

    class Attribute:
        @overload
        def __init__(
            self,
            name: str,
            type: OpSchema.AttrType,
            description: str = "",
            *,
            required: bool = True,
        ) -> None: ...
        @overload
        def __init__(
            self,
            name: str,
            default_value: AttributeProto,
            description: str = "",
        ) -> None: ...
        @property
        def name(self) -> str: ...
        @property
        def description(self) -> str: ...
        @property
        def type(self) -> OpSchema.AttrType: ...
        @property
        def default_value(self) -> AttributeProto: ...
        @property
        def required(self) -> bool: ...

    class SupportType(int):
        COMMON: OpSchema.SupportType = ...
        EXPERIMENTAL: OpSchema.SupportType = ...

    class UseType:
        DEFAULT: OpSchema.UseType = ...
        CONSUME_ALLOWED: OpSchema.UseType = ...
        CONSUME_ENFORCED: OpSchema.UseType = ...

@overload
def has_schema(op_type: str, domain: str = "") -> bool: ...
@overload
def has_schema(op_type: str, max_inclusive_version: int, domain: str = "") -> bool: ...
def schema_version_map() -> dict[str, tuple[int, int]]: ...
@overload
def get_schema(
    op_type: str, max_inclusive_version: int, domain: str = ""
) -> OpSchema: ...
@overload
def get_schema(op_type: str, domain: str = "") -> OpSchema: ...
def get_all_schemas() -> Sequence[OpSchema]: ...
def get_all_schemas_with_history() -> Sequence[OpSchema]: ...
def set_domain_to_version(
    domain: str, min_version: int, max_version: int, last_release_version: int = -1
) -> None: ...
def register_schema(schema: OpSchema) -> None: ...
def deregister_schema(op_type: str, version: int, domain: str) -> None: ...