onnxscript 0.1.0__py3-none-any.whl

onnxscript/_internal/autocast.py

@@ -0,0 +1,250 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence
+
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper
+from onnx.defs import OpSchema
+
+from onnxscript import tensor
+
+if TYPE_CHECKING:
+    from onnxscript import converter
+
+# Conversions from python values to ONNX are used by both the script converter as well
+# as the eager-mode runtime and both need to be consistent. The script converter converts
+# python values into ONNX TensorProto, while the runtime converts python values into
+# ONNXScript runtime's value-representation (based on Tensor).
+
+
+# Utilities to convert a python value to TensorProto (for use by the script converter)
+
+
+def _py_type_to_onnx_type(pytype: type):
+    if pytype is bool:
+        return onnx.TensorProto.BOOL
+    if pytype is int:
+        return onnx.TensorProto.INT64
+    if pytype is float:
+        return onnx.TensorProto.FLOAT
+    if pytype is str:
+        return onnx.TensorProto.STRING
+    raise ValueError(f"Tensor element of type {pytype} not supported")
+
+
+def pyvalue_to_onnx_tensor(tensor_name: str, pyvalue):
+    if isinstance(pyvalue, np.ndarray):
+        return numpy_helper.from_array(pyvalue, tensor_name)
+    if isinstance(pyvalue, list):
+        if len(pyvalue) == 0:
+            raise ValueError("Cannot convert an empty list to tensor")
+        pytype = type(pyvalue[0])
+        if not all(isinstance(e, pytype) for e in pyvalue):
+            raise ValueError(
+                "Cannot convert an list with elements of different types to tensor"
+            )
+        return helper.make_tensor(
+            tensor_name,
+            _py_type_to_onnx_type(pytype),
+            [len(pyvalue)],
+            pyvalue,
+        )
+    onnx_type = _py_type_to_onnx_type(type(pyvalue))
+    if onnx_type is onnx.TensorProto.BOOL:
+        return helper.make_tensor(tensor_name, onnx_type, [], [int(pyvalue)])
+    if onnx_type is onnx.TensorProto.STRING:
+        return helper.make_tensor(tensor_name, onnx_type, [], vals=[pyvalue.encode("utf-8")])
+
+    return helper.make_tensor(tensor_name, onnx_type, [], [pyvalue])
+
+
+_REPEATED_ATTRIBUTE_TYPES = frozenset(
+    {
+        onnx.AttributeProto.FLOATS,
+        onnx.AttributeProto.INTS,
+        onnx.AttributeProto.STRINGS,
+        onnx.AttributeProto.TENSORS,
+        onnx.AttributeProto.GRAPHS,
+        onnx.AttributeProto.SPARSE_TENSORS,
+        onnx.AttributeProto.TYPE_PROTOS,
+    }
+)
+
+
+def pyvalue_to_onnx_attribute(
+    key: str,
+    value: Any,
+    name_generator: Callable[[], str],
+    attr_type: Optional[onnx.AttributeProto.AttributeType] = None,
+) -> onnx.AttributeProto:
+    """Helper function to create an ONNX AttributeProto.
+
+    This is a refinement of onnx.helper.make_attribute that works with ONNX Script
+    conventions for allowed types for attribute-values. In particular, it allows
+    * Empty lists as attribute values, provided the attribute type is specified
+    and is a list type.
+    * Scalar-values like 1.0 as well as lists like [1, -1] to be specified
+    when the attribute type is TensorProto by automatically converting the value
+    into a 0-D or 1-D tensor respectively.
+    """
+    if isinstance(value, list) and not value:
+        # Empty list value:
+        if attr_type is None:
+            raise ValueError("Attribute type must be specified for empty list value.")
+        if attr_type not in _REPEATED_ATTRIBUTE_TYPES:
+            raise ValueError("Empty list value is only allowed for repeated attribute types.")
+        return onnx.AttributeProto(name=key, type=attr_type)
+    elif attr_type == onnx.AttributeProto.TENSOR and not isinstance(value, onnx.TensorProto):
+        return onnx.AttributeProto(
+            name=key, type=attr_type, t=pyvalue_to_onnx_tensor(name_generator(), value)
+        )
+    else:
+        return onnx.helper.make_attribute(key, value)
+
+
+# Utilities to convert python values into onnxscript tensors.
+
+
+def _promotable(x) -> bool:
+    """Checks if a runtime parameter value needs to be promoted into an onnxscript value.
+    This is the runtime-equivalent of the promotion of literal constants into ONNX values
+    in the static converter.
+    """
+    if isinstance(x, (bool, int, float)):
+        return True
+    if isinstance(x, list) and x:
+        # Note: This is meant to handle valid scenarios correctly. No attempt is
+        # made yet to capture all invalid usages in runtime mode.
+        return _promotable(x[0])
+    return False
+
+
+def _get_dtype(pyvalue):
+    """Return np.dtype to use when converting a python value to an onnxscript tensor.
+    Note that int constants are treated as int64, as that is the common type in ONNX
+    for shape/index values.
+    """
+    if isinstance(pyvalue, bool):
+        return np.bool_
+    elif isinstance(pyvalue, int):
+        return np.int64
+    elif isinstance(pyvalue, float):
+        return np.float32
+    elif isinstance(pyvalue, list):
+        if pyvalue:
+            # TODO: What to do about lists with mixed value types, like [1, 2.0]?
+            # Should at least produce an error/warning message.
+            return _get_dtype(pyvalue[0])
+        raise ValueError("Cannot determine target type for empty list")
+    raise TypeError(f"Value of unexpected type {type(pyvalue)}")
+
+
+def cast_pyvalue_to_os_tensor(pyvalue, dtype=None):
+    """Promotes python values into onnxscript tensors.
+    The optional argument dtype specifies the desired np.dtype of the tensor,
+    used only when a non-standard onnxscript-value is promoted into one.
+    """
+    if _promotable(pyvalue):
+        if dtype is None:
+            dtype = _get_dtype(pyvalue)
+        return tensor.Tensor(np.array(pyvalue, dtype=dtype))
+    return pyvalue
+
+
+def cast_inputs(
+    get_type_info: Callable[[Any], Any],
+    cast: Callable[[Any, Any], Any],
+    op_schema: OpSchema | None,
+    args,
+) -> tuple[Any, ...]:
+    """Uses schema specification to support a limited form of auto-casting.
+
+    * Scalars are promoted to tensors.
+    * Further. they are cast to the required type when used in ops with other
+    tensor inputs that are required to be of same type.
+    Thus, in "A+1" or "Add(A, 1)", the value 1 will be converted to the same
+    type as A.
+
+    This is used by the converter in a static-mode, as well as by the eager-mode
+    execution in a dynamic-mode.
+    """
+    if op_schema is None:
+        # Either an error or a custom op.
+        # No checks/casts in this case.
+        return tuple(cast(x, None) for x in args)
+
+    expected_inputs = op_schema.inputs
+    # We make two passes. In the first pass, we identify known type-bindings for
+    # type-variables: eg., {'T1' : np.float32, 'T2' : np.int32}.
+    # In the second pass, we use these bindings to cast scalar-values to
+    # tensors of appropriate types. The two passes are needed to handle cases
+    # like "Add(1, X)" where 1 must be cast to the same type as X.
+    type_bindings: dict[Optional[str], np.dtype] = {}
+    args_typevars: list[tuple[str, Optional[str]]] = []
+    for i, x in enumerate(args):
+        if i < len(expected_inputs):
+            expected = expected_inputs[i]
+        elif expected_inputs[-1].option == OpSchema.FormalParameterOption.Variadic:
+            expected = expected_inputs[-1]
+            if not expected.is_homogeneous:
+                args_typevars.append((x, None))
+                continue
+        else:
+            raise ValueError(
+                f"Number of actual parameters {len(args)} "
+                f"exceeds number of formal parameters {len(expected_inputs)}."
+            )
+        typevar = expected.type_str
+        if "(" not in typevar:
+            # typevar is an identifier, like "T"
+            typeinfo = get_type_info(x)
+            if typeinfo is not None:
+                type_bindings[typevar] = typeinfo
+        args_typevars.append((x, typevar))
+    cast_args = [cast(x, type_bindings.get(typevar)) for x, typevar in args_typevars]
+    return tuple(cast_args)
+
+
+def dynamic_cast_inputs(op_schema: OpSchema, args):
+    """Used for autocast during eager-mode execution."""
+
+    def get_type_info(x):
+        return x.dtype if isinstance(x, tensor.Tensor) else None
+
+    return cast_inputs(get_type_info, cast_pyvalue_to_os_tensor, op_schema, args)
+
+
+def static_cast_inputs(
+    converter_: converter.Converter,
+    op_schema: Optional[OpSchema],
+    args: Sequence[Optional[converter.Variable]],
+) -> tuple[str, ...]:
+    """Used for autocast during script-translation.
+    This is meant to transform expressions like "Add(X, 1)" to "Add(X, CastLike(1, X))"
+    Polymorphic constants (like 0 and 1) are cast to the type of other operands as needed.
+    """
+
+    def get_type_info(x: Optional[converter.Variable]) -> Optional[converter.Variable]:
+        """Returns x back if x can serve as the target-type for a cast (as the second
+        argument of CastLike) and None otherwise. In the expression "Add(X, 1), 1 is
+        castable, while X can serve as the target-type.
+        """
+        return None if x is None or x.is_castable else x
+
+    def cast_like(
+        x: Optional[converter.Variable], y: Optional[converter.Variable]
+    ) -> Optional[str]:
+        if x is None:
+            return None
+        if x.is_castable and y is not None:
+            # Polymorphic constant x is cast to the type of y:
+            x_cast = converter_.generate_unique_name(f"{x.name}_cast")
+            converter_.emit([x_cast], "CastLike", [x.name, y.name])
+            return x_cast
+        return x.name
+
+    return cast_inputs(get_type_info, cast_like, op_schema, args)

onnxscript/_internal/deprecation.py

@@ -0,0 +1,78 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Utility for deprecating APIs."""
+
+# Reference: https://github.com/pytorch/pytorch/blob/aed9bee0413dac190452fbfa9ab2a44b6e6843f5/torch/onnx/_deprecation.py
+
+import functools
+import textwrap
+import warnings
+from typing import Callable, TypeVar
+
+T = TypeVar("T")
+
+
+@functools.lru_cache(maxsize=1024)
+def _warn_once(message: str):
+    """Issue a FutureWarning only once per message."""
+    warnings.warn(message, category=FutureWarning, stacklevel=3)
+
+
+def deprecated(since: str, removed_in: str, instructions: str) -> Callable[[T], T]:
+    """Marks functions as deprecated.
+
+    It will result in a warning when the function is called and a note in the
+    docstring.
+
+    Args:
+        since: The version when the function was first deprecated.
+        removed_in: The version when the function will be removed.
+        instructions: The action users should take.
+
+    Returns:
+        A decorator that can be used to mark functions as deprecated.
+    """
+
+    def decorator(function):
+        @functools.wraps(function)
+        def wrapper(*args, **kwargs):
+            _warn_once(
+                f"'{function.__module__}.{function.__qualname__}' "
+                f"is deprecated in version {since} and will be "
+                f"removed in {removed_in}. Please {instructions}.",
+            )
+            return function(*args, **kwargs)
+
+        # Add a deprecation note to the docstring.
+        docstring = function.__doc__ or ""
+
+        # Add a note to the docstring.
+        deprecation_note = textwrap.dedent(
+            f"""\
+            .. deprecated:: {since}
+                Deprecated and will be removed in version {removed_in}.
+                Please {instructions}.
+            """
+        )
+
+        # Split docstring at first occurrence of newline
+        summary_and_body = docstring.split("\n\n", 1)
+
+        if len(summary_and_body) > 1:
+            summary, body = summary_and_body
+
+            # Dedent the body. We cannot do this with the presence of the summary because
+            # the body contains leading whitespaces when the summary does not.
+            body = textwrap.dedent(body)
+
+            new_docstring_parts = [deprecation_note, "\n\n", summary, body]
+        else:
+            summary = summary_and_body[0]
+
+            new_docstring_parts = [deprecation_note, "\n\n", summary]
+
+        wrapper.__doc__ = "".join(new_docstring_parts)
+
+        return wrapper
+
+    return decorator

onnxscript/_internal/param_manipulation.py

@@ -0,0 +1,148 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""Function for manipulating input parameters of an Op or a OnnxFunction."""
+
+from __future__ import annotations
+
+import collections
+from typing import Any, OrderedDict, Sequence
+
+from onnxscript import values
+
+
+def separate_input_attributes_from_arguments(
+    param_schemas: Sequence[values.ParamSchema],
+    args,
+    kwargs,
+    fill_defaults: bool = True,
+    allow_extra_kwargs: bool = False,
+) -> tuple[list[Any], OrderedDict[str, Any]]:
+    """Separate Python args and kwargs into ONNX inputs and attributes.
+
+    Args:
+        param_schemas: The parameter schemas of an Op or a OnnxFunction.
+        args: The Python positional arguments supplied by the caller.
+        kwargs: The Python keyword arguments supplied by the caller.
+        fill_defaults: Whether to fill the default values for attributes.
+        allow_extra_kwargs: Whether to allow extra keyword arguments.
+            When set to True, extra/unknown arguments will be ignored.
+
+    Returns:
+        A tuple of two elements:
+        - A list of ONNX inputs.
+        - An ordered dictionary of ONNX attribute names and values.
+
+    Raises:
+        TypeError: When allow_extra_kwargs is False and there are unknown kwargs.
+        TypeError: When a required input is not provided.
+    """
+    # args, kwargs and param_schemas should be all in order
+    # user may not specify all inputs or attributes
+
+    all_param_names = {param.name for param in param_schemas}
+    extra_kwargs = set(kwargs).difference(all_param_names)
+    if extra_kwargs and not allow_extra_kwargs:
+        raise TypeError(f"Unexpected keyword arguments '{extra_kwargs}'")
+
+    onnx_inputs = []
+    onnx_attributes = collections.OrderedDict()
+
+    for i, param in enumerate(param_schemas):
+        if param.is_variadic_input:
+            # Exhaust all remaining args
+            onnx_inputs.extend(args[i:])
+            args = []
+            continue
+        if i < len(args):
+            if param.is_input:
+                onnx_inputs.append(args[i])
+            else:
+                onnx_attributes[param.name] = args[i]
+        elif param.name in kwargs:
+            if param.is_input:
+                onnx_inputs.append(kwargs[param.name])
+            else:
+                onnx_attributes[param.name] = kwargs[param.name]
+        elif (
+            param.is_attribute and param.default is not values._EmptyDefault  # pylint: disable=protected-access
+        ):
+            # User did not provide the attribute
+            if fill_defaults:
+                onnx_attributes[param.name] = param.default
+        elif param.required:
+            raise TypeError(f"Required input '{param}' was not provided")
+
+    return onnx_inputs, onnx_attributes
+
+
+def tag_arguments_with_param_schemas(
+    param_schemas: Sequence[values.ParamSchema],
+    args,
+    kwargs,
+    fill_defaults: bool = True,
+    allow_extra_kwargs: bool = False,
+) -> tuple[list[tuple[Any, values.ParamSchema]], dict[str, tuple[Any, values.ParamSchema]]]:
+    """Tag Python args and kwargs with matching ONNX ParamSchema.
+
+    Args:
+        param_schemas: The parameter schemas of an Op or a OnnxFunction.
+        args: The Python positional arguments supplied by the caller.
+        kwargs: The Python keyword arguments supplied by the caller.
+        fill_defaults: Whether to fill the default values for attributes.
+        allow_extra_kwargs: Whether to allow extra keyword arguments.
+            When set to True, extra/unknown arguments will be ignored.
+
+    Returns:
+        A tuple of two elements:
+        - A list of tuple of Python positional argument and ParamSchema.
+        - An ordered dictionary of Python keyword argument names and tuple of argument
+            value and ParamSchema.
+
+    Raises:
+        TypeError: When allow_extra_kwargs is False and there are unknown kwargs.
+        TypeError: When a required input is not provided.
+    """
+    # args, kwargs and param_schemas should be all in order
+    # user may not specify all inputs or attributes
+
+    all_param_names = {param.name for param in param_schemas}
+    extra_kwargs = set(kwargs).difference(all_param_names)
+    if extra_kwargs and not allow_extra_kwargs:
+        raise TypeError(f"Unexpected keyword arguments '{extra_kwargs}'")
+
+    tagged_args: list[tuple[Any, values.ParamSchema]] = []
+    tagged_kwargs: dict[str, tuple[Any, values.ParamSchema]] = {}
+
+    for i, param in enumerate(param_schemas):
+        if param.is_variadic_input:
+            # Exhaust all remaining args
+            tagged_args.extend((arg, param) for arg in args[i:])
+            args = []
+            continue
+        if i < len(args):
+            tagged_args.append((args[i], param))
+        elif param.name in kwargs:
+            tagged_kwargs[param.name] = (kwargs[param.name], param)
+        elif param.default is not values._EmptyDefault:  # pylint: disable=protected-access
+            # User did not provide the input/attribute
+            if fill_defaults:
+                tagged_kwargs[param.name] = (param.default, param)
+        elif param.required:
+            raise TypeError(f"Required input/attribute '{param}' was not provided")
+
+    return tagged_args, tagged_kwargs
+
+
+def turn_to_kwargs_to_avoid_ordering(
+    param_schemas: Sequence[values.ParamSchema],
+    inputs: list[Any],
+    attributes: dict[str, Any],
+) -> dict[str, Any]:
+    """Return the inputs and attributes to the order of the function signature."""
+    for idx, param in enumerate(param_schemas):
+        if param.name not in attributes:
+            if param.is_variadic_input:
+                attributes[param.name] = inputs[idx:]
+            elif inputs:
+                attributes[param.name] = inputs.pop(0)
+    return attributes

onnxscript/_internal/runtime_typing.py

@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+"""An internal wrapper for the beartype library.
+
+Decorate a function with `@runtime_typing.checked` to enable runtime
+type checking. The decorator is a no-op when the `beartype` library is not
+installed.
+"""
+
+import typing
+import warnings
+
+__all__ = [
+    "checked",
+]
+
+T = typing.TypeVar("T", bound=typing.Callable[..., typing.Any])
+
+try:
+    from beartype import beartype as _beartype_decorator
+    from beartype import roar as _roar
+
+    checked = typing.cast(typing.Callable[[T], T], _beartype_decorator)
+
+    # Beartype warns when we import from typing because the types are deprecated
+    # in Python 3.9. But there will be a long time until we can move to using
+    # the native container types for type annotations (when 3.9 is the lowest
+    # supported version). So we silence the warning.
+    warnings.filterwarnings(
+        "ignore",
+        category=_roar.BeartypeDecorHintPep585DeprecationWarning,
+    )
+except ImportError:
+
+    def checked(func: T) -> T:  # type: ignore[no-redef]
+        return func
+
+except Exception as e:  # pylint: disable=broad-exception-caught
+    # Warn errors that are not import errors (unexpected).
+    warnings.warn(f"{e}", stacklevel=2)
+
+    def checked(func: T) -> T:  # type: ignore[no-redef]
+        return func

onnxscript/_internal/utils.py

@@ -0,0 +1,99 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import numbers
+from typing import Optional, Sequence
+
+import numpy as np
+import onnx
+import onnx.helper
+
+from onnxscript import tensor
+
+
+def external_tensor(
+    name: str,
+    data_type: int,
+    dims: Sequence[int],
+    location: str,
+    offset: Optional[int] = None,
+    length: Optional[int] = None,
+    checksum: Optional[str] = None,
+    basepath: Optional[str] = None,
+) -> onnx.TensorProto:
+    """Create a TensorProto referencing externally stored tensor-data.
+
+    Args:
+        name: name of the tensor
+        data_type: data type of tensor element
+        dims: shape of the tensor
+        location: location of the external file (relative path)
+        offset: offset in the file where the tensor-data starts
+        length: number of bytes containing the data
+        checksum: SHA1 digest of the file
+        basepath: basepath combined with location to form the full path
+
+    Returns:
+        TensorProto
+
+    See https://github.com/onnx/onnx/blob/main/docs/ExternalData.md for more details.
+    """
+    tensor_proto = onnx.TensorProto()
+    tensor_proto.name = name
+    tensor_proto.data_type = data_type
+    tensor_proto.dims.extend(dims)
+    tensor_proto.data_location = onnx.TensorProto.EXTERNAL
+
+    def add(k, v):
+        entry = tensor_proto.external_data.add()
+        entry.key = k
+        entry.value = str(v)
+
+    add("location", location)
+    if offset is not None:
+        add("offset", int(offset))
+    if length is not None:
+        add("length", int(length))
+    if checksum is not None:
+        add("checksum", checksum)
+    if basepath is not None:
+        add("basepath", basepath)
+    return tensor_proto
+
+
+def value_to_type_proto(val):
+    """Return the ONNX type of a python-value."""
+    if isinstance(val, (np.ndarray, tensor.Tensor)):
+        elem_type = onnx.helper.np_dtype_to_tensor_dtype(val.dtype)
+        shape = val.shape
+        return onnx.helper.make_tensor_type_proto(elem_type, shape)
+    if isinstance(val, int):
+        return onnx.helper.make_tensor_type_proto(onnx.TensorProto.INT32, [])
+    if isinstance(val, (float, np.float32)):
+        return onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, [])
+    if isinstance(val, list):
+        if len(val) > 0:
+            return onnx.helper.make_sequence_type_proto(value_to_type_proto(val[0]))
+        # Edge-case. Cannot determine a suitable ONNX type for an empty list.
+        # Should be using a typed-value instead.
+        # Treated as a sequence of tensors of float-type.
+        return onnx.helper.make_sequence_type_proto(
+            onnx.helper.make_tensor_type_proto(onnx.TensorProto.FLOAT, None)
+        )
+    if isinstance(val, numbers.Number):
+        nparray = np.array(val)
+        elem_type = onnx.helper.np_dtype_to_tensor_dtype(nparray.dtype)
+        return onnx.helper.make_tensor_type_proto(elem_type, [])
+    raise ValueError(f"Value of type {type(val)} is invalid as an ONNX input/output.")
+
+
+def values_to_value_infos(name_values):
+    """Create a list of ValueInfoProto from a list of (name, value) pairs,
+    skipping any None values.
+    """
+    return [
+        onnx.helper.make_value_info(name, value_to_type_proto(val))
+        for (name, val) in name_values
+        if val is not None
+    ]