onnxscript 0.1.0.dev20231130__tar.gz → 0.1.0.dev20231201__tar.gz

{onnxscript-0.1.0.dev20231130/onnxscript.egg-info → onnxscript-0.1.0.dev20231201}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxscript
-Version: 0.1.0.dev20231130
+Version: 0.1.0.dev20231201
 Summary: Naturally author ONNX functions and models using a subset of Python
 Home-page: https://onnxscript.ai/
 Author-email: Microsoft Corporation <onnx@microsoft.com>
@@ -111,9 +111,8 @@ pytest onnxscript
 import onnx
 
 # We use ONNX opset 15 to define the function below.
-from onnxscript import FLOAT
+from onnxscript import FLOAT, script
 from onnxscript import opset15 as op
-from onnxscript import script
 
 
 # We use the script decorator to indicate that

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201}/README.md

@@ -65,9 +65,8 @@ pytest onnxscript
 import onnx
 
 # We use ONNX opset 15 to define the function below.
-from onnxscript import FLOAT
+from onnxscript import FLOAT, script
 from onnxscript import opset15 as op
-from onnxscript import script
 
 
 # We use the script decorator to indicate that

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201}/onnxscript/__init__.py

@@ -58,7 +58,7 @@ from .onnx_types import (
 
 # isort: on
 
-from ._internal.utils import external_tensor, proto2text
+from ._internal.utils import external_tensor
 from .values import OnnxFunction, TracedOnnxFunction
 
 try:
@@ -73,7 +73,6 @@ __all__ = [
     "OnnxFunction",
     "TracedOnnxFunction",
     "proto2python",
-    "proto2text",
     "external_tensor",
     "graph",
     "BFLOAT16",

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201}/onnxscript/_internal/utils.py

@@ -14,16 +14,6 @@ from onnx import FunctionProto, ModelProto, TensorProto, ValueInfoProto
 
 from onnxscript import tensor
 
-# print utility unavailable in ONNX 1.12 or earlier:
-# pylint: disable=unused-import, ungrouped-imports
-try:
-    from onnx.printer import to_text as proto2text
-except ImportError:
-
-    def proto2text(_: Any) -> str:  # type: ignore[misc]
-        return "<print utility unavailable>"
-
-
 # pylint: enable=unused-import, ungrouped-imports
 
 

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201}/onnxscript/backend/onnx_export_test.py

@@ -11,6 +11,7 @@ import re
 import unittest
 from typing import Pattern
 
+import onnx
 import onnxruntime as ort
 import parameterized
 from onnxruntime.capi import onnxruntime_pybind11_state
@@ -273,7 +274,7 @@ class TestOnnxBackEnd(unittest.TestCase):
         except Exception as e:
             raise AssertionError(
                 f"Unable to load onnx for test {backend_test.name!r}.\n"
-                f"{onnxscript.proto2text(proto)}\n"
+                f"{onnx.printer.to_text(proto)}\n"
                 f"-----\n"
                 f"{backend_test.onnx_model}"
             ) from e
@@ -298,7 +299,7 @@ class TestOnnxBackEnd(unittest.TestCase):
             except Exception as e:
                 raise AssertionError(
                     f"Unable to run test {backend_test.name!r} after conversion.\n"
-                    f"{onnxscript.proto2text(proto)}"
+                    f"{onnx.printer.to_text(proto)}"
                 ) from e
 
         backend_test.run(_load_function, _run_function)

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201}/onnxscript/evaluator.py

@@ -480,7 +480,7 @@ def _call_ort(
         raise EagerModeError(
             f"Unable to create onnxruntime InferenceSession "
             f"for executing {schema.domain}.{schema.name} op "
-            f"with onnx model\n{utils.proto2text(model)}"
+            f"with onnx model\n{onnx.printer.to_text(model)}"
         ) from e
 
     try:

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201}/onnxscript/function_libs/torch_lib/graph_building.py

@@ -1029,7 +1029,7 @@ class TorchScriptGraph:
             warnings.warn(f"ONNX model is invalid: {e}", stacklevel=1)
             logging.debug(
                 "ONNX model:\n%s\n\nTorchScript graph:\n%s",
-                onnxscript.proto2text(onnx_model),
+                onnx.printer.to_text(onnx_model),
                 self.torch_graph,
             )
         return onnx_model

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201}/onnxscript/function_libs/torch_lib/ops/core.py

@@ -4505,40 +4505,30 @@ def aten_linear_backward(
 
 @torch_op("aten::linspace", trace_only=True)
 def aten_linspace(
-    start: TRealUnlessFloat16OrInt8, end: TRealUnlessFloat16OrInt8, steps: int, dtype: int = -1
+    start: TFloat, end: TFloat, steps: int, dtype: int = FLOAT.dtype
 ) -> TensorType:
     """linspace(Scalar start, Scalar end, int steps, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor"""
 
-    if dtype == -1:
-        zero = op.CastLike(0.0, steps)
-        one = op.CastLike(1.0, steps)
-    elif _range_supported(dtype):
-        zero = op.Cast(0, to=dtype)
-        one = op.Cast(1, to=dtype)
-        start = op.Cast(start, to=dtype)
-        end = op.Cast(end, to=dtype)
-        steps = op.Cast(steps, to=dtype)
-    else:
-        # Cast input to float if dtype is not supported by Range,
-        # because the input dtype may be e.g. bfloat16 / int8 etc.
-        # which Range does not support. The output type is ensured because the output
-        # is casted to the specified dtype.
-        zero = op.Cast(0.0, to=FLOAT.dtype)
-        one = op.Cast(1.0, to=FLOAT.dtype)
-        start = op.Cast(start, to=FLOAT.dtype)
-        end = op.Cast(end, to=FLOAT.dtype)
-        steps = op.Cast(steps, to=FLOAT.dtype)
-
-    range_tensor = op.Range(zero, steps, one)
-
-    start = op.CastLike(start, end)
-    step = op.Div(
-        op.Sub(end, start),
-        op.Sub(steps, one),
+    # Reference: https://github.com/pytorch/pytorch/blob/b35ca2cb941b5ba90858322810ca85c31e4541fd/torch/_refs/__init__.py#L4896
+    if steps == 0:
+        return aten_full(op.Constant(value_ints=[0]), 0.0, dtype=dtype)
+    if steps == 1:
+        return aten_full(op.Constant(value_ints=[steps]), start, dtype=dtype)
+
+    rg = aten_arange_start(0, steps, dtype=dtype)
+    start = op.Cast(start, to=dtype)
+    end = op.Cast(end, to=dtype)
+    steps_float = op.Cast(steps, to=dtype)
+    one = op.Cast(1.0, to=dtype)
+    two = op.Cast(2.0, to=dtype)
+    steps_minus_1 = op.Cast(steps - 1, to=dtype)
+    step = op.Div(op.Sub(end, start), steps_minus_1)
+    return op.Where(
+        rg < op.Div(steps_float, two),
+        start + step * rg,
+        end - step * (steps_float - one - rg),
     )
 
-    return op.Add(op.Mul(range_tensor, step), start)
-
 
 @torch_op("aten::log")
 def aten_log(self: TFloatOrBFloat16) -> TFloatOrBFloat16:
@@ -6925,6 +6915,37 @@ def aten_roll(self: TTensor, shifts: INT64, dims: Sequence[int] = ()) -> TTensor
             return result
 
 
+@torch_op("aten::roll", trace_only=True, complex=True)
+def aten_roll_complex(self: TTensor, shifts: INT64, dims: Sequence[int] = ()) -> TTensor:
+    """roll(Tensor self, int[1] shifts, int[1] dims=[]) -> Tensor"""
+
+    self_rank = len(self.shape)
+    if self_rank == 1:
+        return self
+
+    if self.shape[0] == 0:  # empty tensor
+        return self
+
+    self_real = op.Slice(self, [0], [1], axes=[-1])
+    self_imag = op.Slice(self, [1], [2], axes=[-1])
+    if not dims:
+        # assert isinstance(shifts, int)
+        shift_real = _aten_roll_shift_no_dim_onnx(self_real, shifts)
+        shift_imag = _aten_roll_shift_no_dim_onnx(self_imag, shifts)
+
+        result = op.Concat(shift_real, shift_imag, axis=-1)
+
+    else:
+        # assert len(shifts) == len(dims), but shifts is a tensor, dims is a list
+        for i, dim in enumerate(dims):
+            shift = op.Gather(shifts, i, axis=0)
+            self_real = _aten_roll_shift_and_dim_onnx(self_real, shift, dim)
+            self_imag = _aten_roll_shift_and_dim_onnx(self_imag, shift, dim)
+
+        result = op.Concat(self_real, self_imag, axis=-1)
+    return result
+
+
 @torch_op("aten::roll", private=True)
 def _aten_roll_shift_no_dim_onnx(self: TTensor, shift: INT64) -> TTensor:
     neg_1 = op.Constant(value_ints=[-1])
@@ -8266,10 +8287,47 @@ def aten_vander(
     raise NotImplementedError()
 
 
-def aten_var(self: TensorType, unbiased: bool = True) -> TensorType:
+@torch_op("aten::var", trace_only=True)
+def aten_var(self: TReal, unbiased: Optional[bool] = True) -> TReal:
     """var(Tensor self, bool unbiased=True) -> Tensor"""
 
-    raise NotImplementedError()
+    # Assume bool(True) and int(1) are same in ONNX, so pass "unbiased" directly as "correction"
+    # If not this case, should be explicitly set correction value according to unbiased value
+    return _aten_var_onnx(self, correction=float(unbiased), keepdim=False)
+
+
+@torch_op("aten::var.dim", trace_only=True)
+def aten_var_dim(
+    self: TReal, dim: int, unbiased: Optional[bool] = True, keepdim: Optional[bool] = False
+) -> TReal:
+    """var(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False) -> Tensor"""
+
+    if isinstance(dim, int):
+        dim = (dim,)
+    dim_tensor = op.Constant(value_ints=dim)
+    return _aten_var_dim_onnx(self, dim_tensor, correction=float(unbiased), keepdim=keepdim)
+
+
+@torch_op("aten::var.correction", trace_only=True)
+def aten_var_correction(
+    self: TReal,
+    dim: Optional[int] = None,
+    correction: Optional[float] = None,
+    keepdim: bool = False,
+) -> TReal:
+    """var.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> Tensor"""
+
+    if correction is None:
+        correction = 1.0
+
+    if dim is None:
+        var = _aten_var_onnx(self, correction, keepdim)
+    else:
+        if isinstance(dim, int):
+            dim = (dim,)
+        dim_tensor = op.Constant(value_ints=dim)
+        var = _aten_var_dim_onnx(self, dim_tensor, correction, keepdim)
+    return var
 
 
 @torch_op("aten::var_mean", trace_only=True)
@@ -8355,12 +8413,50 @@ def _aten_var_mean_dim_onnx(
         dim_size = op.Gather(self_shape, dim, axis=0)
         numel_float = op.CastLike(op.ReduceProd(dim_size, keepdims=False), self)
         mul = op.Mul(var, numel_float)
-        sub = op.Sub(numel_float, correction)
+        sub = op.Sub(numel_float, op.CastLike(correction, self))
         var = op.Div(mul, sub)
 
     return var, mean
 
 
+@torch_op("aten::var", private=True, traceable=True)
+def _aten_var_onnx(self: TReal, correction: float, keepdim: bool = False) -> TReal:
+    mean = op.ReduceMean(self, keepdims=keepdim)
+    sub_mean = op.Sub(self, mean)
+    sqr_mean = op.Mul(sub_mean, sub_mean)
+    var = op.ReduceMean(sqr_mean, keepdims=keepdim)
+    # Adjust var according to correction value
+    if correction > 0.0:
+        self_shape = op.Shape(self)
+        numel_float = op.CastLike(op.ReduceProd(self_shape, keepdims=False), self)
+        mul = op.Mul(var, numel_float)
+        sub = op.Sub(numel_float, op.CastLike(correction, self))
+        var = op.Div(mul, sub)
+
+    return var
+
+
+@torch_op("aten::var.dim", private=True, traceable=True)
+def _aten_var_dim_onnx(
+    self: TReal, dim: INT64, correction: float, keepdim: bool = False
+) -> TReal:
+    dim = op.Reshape(dim, op.Constant(value_ints=[-1]))
+    # Computer mean and var
+    sub_mean = op.Sub(self, op.ReduceMean(self, dim, keepdims=True))
+    sqr_mean = op.Mul(sub_mean, sub_mean)
+    var = op.ReduceMean(sqr_mean, dim, keepdims=keepdim)
+    # Adjust var according to correction value
+    if correction > 0.0:
+        self_shape = op.Shape(self)
+        dim_size = op.Gather(self_shape, dim, axis=0)
+        numel_float = op.CastLike(op.ReduceProd(dim_size, keepdims=False), self)
+        mul = op.Mul(var, numel_float)
+        sub = op.Sub(numel_float, op.CastLike(correction, self))
+        var = op.Div(mul, sub)
+
+    return var
+
+
 def aten_vdot(self: TensorType, other: TensorType) -> TensorType:
     """vdot(Tensor self, Tensor other) -> Tensor"""
 

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201}/onnxscript/tests/function_libs/torch_lib/ops_test_common.py

@@ -359,7 +359,7 @@ def _format_model_and_input_information(onnx_model, inputs):
         f"Inputs:\n"
         f"{pprint.pformat(inputs)}\n"
         f"Model:\n"
-        f"{onnxscript.proto2text(onnx_model)}"
+        f"{onnx.printer.to_text(onnx_model)}"
     )
 
 
@@ -526,7 +526,7 @@ def graph_executor(
             onnx.checker.check_model(onnx_model, full_check=True)
         except (onnx.checker.ValidationError, onnx.shape_inference.InferenceError) as e:
             raise AssertionError(
-                f"ONNX model is invalid. Model:\n{onnxscript.proto2text(onnx_model)}"
+                f"ONNX model is invalid. Model:\n{onnx.printer.to_text(onnx_model)}"
             ) from e
 
         try:

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201}/onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -876,6 +876,7 @@ TESTED_TORCHLIB_OPS: tuple[TorchLibOpInfo, ...] = (
         "linspace",
         core_ops.aten_linspace,
         trace_only=True,
+        tolerance={torch.float16: (2e-2, 2e-3)},
     )
     .xfail(
         dtypes=(torch.int64, torch.int32),
@@ -883,17 +884,9 @@ TESTED_TORCHLIB_OPS: tuple[TorchLibOpInfo, ...] = (
     )
     .xfail(
         variant_name="tensor_overload",
-        dtypes=(torch.int64, torch.int32, torch.float16),
+        dtypes=(torch.int64, torch.int32),
         reason="fixme: Results do not match with PyTorch. https://github.com/microsoft/onnxscript/issues/854",
         enabled_if=not version_utils.torch_older_than("2.2"),
-    )
-    .xfail(
-        dtypes=(torch.float16,),
-        reason="op 'Range' doesn't support float16.",
-    )
-    .skip(
-        matcher=lambda sample: len(sample.args) > 1 and sample.args[1] == 1,
-        reason="aten::linspace with steps=1 is not supported by its definition.",
     ),
     TorchLibOpInfo("log", core_ops.aten_log),
     TorchLibOpInfo("le", core_ops.aten_le),
@@ -2075,6 +2068,13 @@ TESTED_TORCHLIB_OPS: tuple[TorchLibOpInfo, ...] = (
         trace_only=True,
         input_wrangler=_roll_input_wrangler,
     ),
+    TorchLibOpInfo(
+        "roll",
+        core_ops.aten_roll_complex,
+        input_wrangler=_roll_input_wrangler,
+        trace_only=True,
+        complex=True,
+    ),
     TorchLibOpInfo(
         "scatter_reduce",
         core_ops.aten_scatter_reduce,
@@ -2182,6 +2182,36 @@ TESTED_TORCHLIB_OPS: tuple[TorchLibOpInfo, ...] = (
         matcher=lambda sample: len(sample.args) > 0 or "correction" not in sample.kwargs,
         reason="this Aten overload only support when correction attribute exists",
     ),
+    TorchLibOpInfo(
+        "var",
+        core_ops.aten_var,
+        trace_only=True,
+    ).xfail(
+        # kwargs must be empty
+        matcher=lambda sample: len(sample.kwargs) > 0,
+        reason="this Aten overload only support input[0]=tensor and input[1]=bool as input without any kwargs",
+    ),
+    TorchLibOpInfo(
+        "var_dim",
+        core_ops.aten_var_dim,
+        trace_only=True,
+    ).xfail(
+        # kwargs["dim"] must exist, kwargs["correction"] must not exist
+        matcher=lambda sample: not (
+            sample.kwargs.get("dim", None) is not None
+            and sample.kwargs.get("correction", None) is None
+        ),
+        reason="this Aten overload only support with 'dim' argument and without 'correction' argument",
+    ),
+    TorchLibOpInfo(
+        "var_correction",
+        core_ops.aten_var_correction,
+        trace_only=True,
+    ).skip(
+        # Don't accept input[1]=bool and 'correction' must be in kwargs
+        matcher=lambda sample: len(sample.args) > 0 or "correction" not in sample.kwargs,
+        reason="this Aten overload only support when correction attribute exists",
+    ),
     TorchLibOpInfo("zeros_like", core_ops.aten_zeros_like, trace_only=True),
 )
 
@@ -2279,6 +2309,7 @@ ops_test_common.duplicate_opinfo(OPS_DB, "ops.aten._softmax", ("ops.aten._softma
 ops_test_common.duplicate_opinfo(OPS_DB, "round", ("round_decimals",))
 ops_test_common.duplicate_opinfo(OPS_DB, "squeeze", ("squeeze_dim",))
 ops_test_common.duplicate_opinfo(OPS_DB, "var_mean", ("var_mean_dim", "var_mean_correction"))
+ops_test_common.duplicate_opinfo(OPS_DB, "var", ("var_dim", "var_correction"))
 ops_test_common.duplicate_opinfo(OPS_DB, "view_as_complex", ("view_as_complex_copy",))
 ops_test_common.duplicate_opinfo(OPS_DB, "view_as_real", ("view_as_real_copy",))
 

{onnxscript-0.1.0.dev20231130 → onnxscript-0.1.0.dev20231201/onnxscript.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxscript
-Version: 0.1.0.dev20231130
+Version: 0.1.0.dev20231201
 Summary: Naturally author ONNX functions and models using a subset of Python
 Home-page: https://onnxscript.ai/
 Author-email: Microsoft Corporation <onnx@microsoft.com>
@@ -111,9 +111,8 @@ pytest onnxscript
 import onnx
 
 # We use ONNX opset 15 to define the function below.
-from onnxscript import FLOAT
+from onnxscript import FLOAT, script
 from onnxscript import opset15 as op
-from onnxscript import script
 
 
 # We use the script decorator to indicate that