From db30dbb0489b47621a0439bbcb180be0919e922a Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 8 Oct 2024 10:14:15 -0700 Subject: [PATCH 1/6] chore(deps): bump ruff from 0.6.8 to 0.6.9 in /requirements/lintrunner (#1891) Bumps [ruff](https://github.com/astral-sh/ruff) from 0.6.8 to 0.6.9.
Release notes

Sourced from ruff's releases.

0.6.9

Release Notes

Preview features

Rule changes

Configuration

Bug fixes

Documentation

Install ruff 0.6.9

Install prebuilt binaries via shell script

curl --proto '=https' --tlsv1.2 -LsSf
https://github.com/astral-sh/ruff/releases/download/0.6.9/ruff-installer.sh
| sh

Install prebuilt binaries via powershell script

powershell -ExecutionPolicy ByPass -c "irm
https://github.com/astral-sh/ruff/releases/download/0.6.9/ruff-installer.ps1
| iex"

Download ruff 0.6.9

File Platform Checksum
ruff-aarch64-apple-darwin.tar.gz Apple Silicon macOS checksum
ruff-x86_64-apple-darwin.tar.gz Intel macOS checksum
ruff-aarch64-pc-windows-msvc.zip ARM64 Windows checksum

... (truncated)

Changelog

Sourced from ruff's changelog.

0.6.9

Preview features

Rule changes

Configuration

Bug fixes

Documentation

Commits

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=ruff&package-manager=pip&previous-version=0.6.8&new-version=0.6.9)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot merge` will merge this PR after your CI passes on it - `@dependabot squash and merge` will squash and merge this PR after your CI passes on it - `@dependabot cancel merge` will cancel a previously requested merge and block automerging - `@dependabot reopen` will reopen this PR if it is closed - `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- requirements/lintrunner/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/lintrunner/requirements.txt b/requirements/lintrunner/requirements.txt index c1643f100..5aa076eb2 100644 --- a/requirements/lintrunner/requirements.txt +++ b/requirements/lintrunner/requirements.txt @@ -1,7 +1,7 @@ # This file is auto updated by dependabot lintrunner-adapters>=0.8.0 # RUFF, RUFF-FIX -ruff==0.6.8 +ruff==0.6.9 # MYPY mypy==1.10.1 types-PyYAML==6.0.12.20240808 From 3be8fc482bc445b6eee4a83205bee5a73279bf1a Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 8 Oct 2024 11:49:08 -0700 Subject: [PATCH 2/6] [torchlib] Include bfloat16 as part of the float types (#1894) Since onnx in opset 20 or so enabled bfloat16 for most relevant ops, we are just going to include allow them in torchlib (even though it is opset18 for now) to unblock bfloat16 model export. --- .../function_libs/torch_lib/ops/core.py | 53 +++++++++---------- onnxscript/function_libs/torch_lib/ops/nn.py | 11 ++-- .../function_libs/torch_lib/ops/special.py | 16 +++--- .../function_libs/torch_lib/tensor_typing.py | 3 +- 4 files changed, 38 insertions(+), 45 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index f41ff1c3e..1fc73a220 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -39,7 +39,6 @@ RealType, TFloat, TFloatHighPrecision, - TFloatOrBFloat16, TInt, TReal, TRealOrUInt8, @@ -3564,14 +3563,14 @@ def aten_flipud(self: TensorType) -> TensorType: @torch_op("aten::floor", traceable=True) -def aten_floor(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_floor(self: TFloat) -> TFloat: """floor(Tensor self) -> Tensor""" return op.Floor(self) @torch_op("math::floor", traceable=True) -def python_math_floor(self: TFloatOrBFloat16) -> TInt: +def python_math_floor(self: TFloat) -> TInt: """floor(Tensor self) -> Tensor""" floor = op.Floor(self) return op.Cast(floor, to=INT64.dtype) @@ -4533,7 +4532,7 @@ def aten_isfinite(self: TFloatHighPrecision) -> BOOL: @torch_op("aten::isinf") -def aten_isinf(self: TFloatOrBFloat16) -> BOOL: +def aten_isinf(self: TFloat) -> BOOL: """isinf(Tensor self) -> Tensor""" # Added Cast inside the function so it can support all real dtypes naturally @@ -4542,14 +4541,14 @@ def aten_isinf(self: TFloatOrBFloat16) -> BOOL: @torch_op("aten::isnan") -def aten_isnan(self: TFloatOrBFloat16) -> BOOL: +def aten_isnan(self: TFloat) -> BOOL: """isnan(Tensor self) -> Tensor""" return op.IsNaN(self) @torch_op("aten::isneginf") -def aten_isneginf(self: TFloatOrBFloat16) -> BOOL: +def aten_isneginf(self: TFloat) -> BOOL: """isneginf(Tensor self) -> Tensor""" # Added Cast inside the function so it can support all real dtypes naturally @@ -4558,7 +4557,7 @@ def aten_isneginf(self: TFloatOrBFloat16) -> BOOL: @torch_op("aten::isposinf") -def aten_isposinf(self: TFloatOrBFloat16) -> BOOL: +def aten_isposinf(self: TFloat) -> BOOL: """isposinf(Tensor self) -> Tensor""" # Added Cast inside the function so it can support all real dtypes naturally @@ -4778,42 +4777,42 @@ def aten_linspace( @torch_op("aten::log", traceable=True) -def aten_log(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_log(self: TFloat) -> TFloat: """log(Tensor self) -> Tensor""" return op.Log(self) @torch_op("aten::log10", traceable=True) -def aten_log10(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_log10(self: TFloat) -> TFloat: """log10(Tensor self) -> Tensor""" return op.Div(op.Log(self), op.CastLike(op.Log(10.0), self)) @torch_op("aten::log1p") -def aten_log1p(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_log1p(self: TFloat) -> TFloat: """log1p(Tensor self) -> Tensor""" return op.Log(op.Add(self, 1.0)) @torch_op("aten::log2", traceable=True) -def aten_log2(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_log2(self: TFloat) -> TFloat: """log2(Tensor self) -> Tensor""" return op.Div(op.Log(self), op.CastLike(op.Log(2.0), self)) @torch_op("aten::logaddexp", traceable=True) -def aten_logaddexp(self: TFloatOrBFloat16, other: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_logaddexp(self: TFloat, other: TFloat) -> TFloat: """logaddexp(Tensor self, Tensor other) -> Tensor""" return op.Log(op.Add(op.Exp(self), op.Exp(other))) @torch_op("aten::logaddexp2", traceable=True) -def aten_logaddexp2(self: TFloatOrBFloat16, other: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_logaddexp2(self: TFloat, other: TFloat) -> TFloat: """logaddexp2(Tensor self, Tensor other) -> Tensor""" two = op.CastLike(2.0, self) summation = op.Add(op.Pow(two, self), op.Pow(two, other)) @@ -4822,7 +4821,7 @@ def aten_logaddexp2(self: TFloatOrBFloat16, other: TFloatOrBFloat16) -> TFloatOr @torch_op("aten::logcumsumexp", traceable=True) -def aten_logcumsumexp(self: TFloatOrBFloat16, dim: int) -> TFloatOrBFloat16: +def aten_logcumsumexp(self: TFloat, dim: int) -> TFloat: """logcumsumexp(Tensor self, int dim) -> Tensor""" if IsScalar(self): @@ -4908,12 +4907,12 @@ def aten_logical_xor(self: BOOL, other: BOOL) -> BOOL: @torch_op("aten::logit", private=True) -def _aten_logit_onnx(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def _aten_logit_onnx(self: TFloat) -> TFloat: return op.Log(op.Div(self, op.Sub(1.0, self))) @torch_op("aten::logit", private=True) -def _aten_logit_clamp_onnx(self: TFloatOrBFloat16, eps: float) -> TFloatOrBFloat16: +def _aten_logit_clamp_onnx(self: TFloat, eps: float) -> TFloat: eps = op.CastLike(eps, self) one = op.CastLike(1.0, self) temporary_self = op.Where(self <= one - eps, self, one - eps) @@ -4923,7 +4922,7 @@ def _aten_logit_clamp_onnx(self: TFloatOrBFloat16, eps: float) -> TFloatOrBFloat @torch_op("aten::logit", trace_only=True) -def aten_logit(self: TFloatOrBFloat16, eps: Optional[float] = None) -> TFloatOrBFloat16: +def aten_logit(self: TFloat, eps: Optional[float] = None) -> TFloat: """logit(Tensor self, float? eps=None) -> Tensor""" if eps is None: return _aten_logit_onnx(self) @@ -6041,9 +6040,7 @@ def aten_native_channel_shuffle(self: TensorType, groups: int) -> TensorType: @torch_op("aten::native_dropout", trace_only=True) -def aten_native_dropout( - input: TFloatOrBFloat16, p: float, train: bool = True -) -> Tuple[TFloatOrBFloat16, BOOL]: +def aten_native_dropout(input: TFloat, p: float, train: bool = True) -> Tuple[TFloat, BOOL]: """native_dropout(Tensor input, float p, bool? train) -> (Tensor, Tensor)""" result, mask = op.Dropout(input, p, train) @@ -7055,7 +7052,7 @@ def aten_real(self: TensorType) -> TensorType: @torch_op("aten::reciprocal") -def aten_reciprocal(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_reciprocal(self: TFloat) -> TFloat: """reciprocal(Tensor self) -> Tensor""" return op.Reciprocal(self) @@ -7074,7 +7071,7 @@ def aten_refine_names(self: TensorType, names: Sequence[str]) -> TensorType: @torch_op(("aten::remainder.Tensor", "aten::remainder.Scalar")) -def aten_remainder(self: TFloatOrBFloat16, other: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_remainder(self: TFloat, other: TFloat) -> TFloat: """remainder.Tensor(Tensor self, Tensor other) -> Tensor""" # TODO(justinchuby): Improve fp16 precision by following the logic in @@ -7355,7 +7352,7 @@ def aten_rrelu( @torch_op("aten::rsqrt", traceable=True) -def aten_rsqrt(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_rsqrt(self: TFloat) -> TFloat: """rsqrt(Tensor self) -> Tensor""" return op.Reciprocal(op.Sqrt(self)) @@ -7562,7 +7559,7 @@ def aten_sgn(self: TensorType) -> TensorType: @torch_op("aten::sigmoid", traceable=True) -def aten_sigmoid(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_sigmoid(self: TFloat) -> TFloat: """sigmoid(Tensor self) -> Tensor""" return op.Sigmoid(self) @@ -7724,7 +7721,7 @@ def aten_smm(self: TensorType, mat2: TensorType) -> TensorType: @torch_op(("aten::softmax.int", "aten::special_softmax"), trace_only=True) -def aten_softmax(self: TFloatOrBFloat16, dim: int, dtype: int = -1) -> TFloatOrBFloat16: +def aten_softmax(self: TFloat, dim: int, dtype: int = -1) -> TFloat: """softmax(Tensor self, int dim, ScalarType? dtype=None) -> Tensor""" self_is_scalar = IsScalar(self) @@ -7741,7 +7738,7 @@ def aten_softmax(self: TFloatOrBFloat16, dim: int, dtype: int = -1) -> TFloatOrB @torch_op(("aten::softmax.int", "aten::special_softmax"), traceable=True) -def aten_softmax_no_dtype(self: TFloatOrBFloat16, dim: int) -> TFloatOrBFloat16: +def aten_softmax_no_dtype(self: TFloat, dim: int) -> TFloat: """softmax(Tensor self, int dim, ScalarType? dtype=None) -> Tensor""" self_is_scalar = IsScalar(self) @@ -7812,7 +7809,7 @@ def aten_split_with_sizes_copy( @torch_op("aten::sqrt", traceable=True) -def aten_sqrt(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_sqrt(self: TFloat) -> TFloat: """sqrt(Tensor self) -> Tensor""" return op.Sqrt(self) @@ -8402,7 +8399,7 @@ def aten_triu_indices(row: int, col: int, offset: int = 0) -> TensorType: @torch_op("aten::trunc") -def aten_trunc(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_trunc(self: TFloat) -> TFloat: """trunc(Tensor self) -> Tensor""" # Reference https://github.com/onnx/onnx/issues/4588#issuecomment-1463970126 diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index 4687e260a..e963050f5 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -25,7 +25,6 @@ from onnxscript.function_libs.torch_lib.tensor_typing import ( IntType, TFloat, - TFloatOrBFloat16, TFloatOrUInt8, TInt, TReal, @@ -364,13 +363,13 @@ def aten_conv_depthwise3d( @torch_op("aten::cross_entropy_loss", traceable=True) def aten_cross_entropy_loss( - self: TFloatOrBFloat16, + self: TFloat, target: IntType, - weight: Optional[TFloatOrBFloat16] = None, + weight: Optional[TFloat] = None, reduction: int = 1, # default is 'mean' ignore_index: int = -100, label_smoothing: float = 0.0, # this was ignored due to ONNX not support -) -> TFloatOrBFloat16: +) -> TFloat: """cross_entropy_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100, float label_smoothing=0.0) -> Tensor""" if reduction == 0: # "none" @@ -812,7 +811,7 @@ def aten_l1_loss(self: TensorType, target: TensorType, reduction: int = 1) -> Te @torch_op("aten::leaky_relu") -def aten_leaky_relu(self: TFloatOrBFloat16, negative_slope: float = 0.01) -> TFloatOrBFloat16: +def aten_leaky_relu(self: TFloat, negative_slope: float = 0.01) -> TFloat: """leaky_relu(Tensor self, Scalar negative_slope=0.01) -> Tensor""" return op.LeakyRelu(self, alpha=negative_slope) @@ -850,7 +849,7 @@ def aten_linear_bias(input: TFloat, weight: TFloat, bias: TFloat) -> TFloat: @torch_op("aten::log_sigmoid") -def aten_log_sigmoid(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_log_sigmoid(self: TFloat) -> TFloat: """log_sigmoid(Tensor self) -> Tensor""" return op.Log(op.Sigmoid(self)) diff --git a/onnxscript/function_libs/torch_lib/ops/special.py b/onnxscript/function_libs/torch_lib/ops/special.py index 6dd9edcd3..c791937b1 100644 --- a/onnxscript/function_libs/torch_lib/ops/special.py +++ b/onnxscript/function_libs/torch_lib/ops/special.py @@ -17,7 +17,7 @@ from onnxscript.function_libs.torch_lib.ops import common as common_ops from onnxscript.function_libs.torch_lib.registration import torch_op -from onnxscript.function_libs.torch_lib.tensor_typing import TFloat, TFloatOrBFloat16 +from onnxscript.function_libs.torch_lib.tensor_typing import TFloat from onnxscript.onnx_opset import opset18 as op from onnxscript.onnx_types import TensorType @@ -92,21 +92,21 @@ def aten_special_entr(self: TensorType) -> TensorType: @torch_op(("aten::erf", "aten::special_erf")) -def aten_special_erf(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_special_erf(self: TFloat) -> TFloat: """erf(Tensor self) -> Tensor""" return op.Erf(self) @torch_op(("aten::erfc", "aten::special_erfc")) -def aten_special_erfc(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_special_erfc(self: TFloat) -> TFloat: """erfc(Tensor self) -> Tensor""" return op.Sub(1, op.Erf(self)) @torch_op("aten::special_erfcx") -def aten_special_erfcx(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_special_erfcx(self: TFloat) -> TFloat: """special_erfcx(Tensor self) -> Tensor""" return op.Mul(op.Exp(op.Pow(self, 2)), op.Sub(1, op.Erf(self))) @@ -131,7 +131,7 @@ def aten_special_expit(self: TensorType) -> TensorType: @torch_op(("aten::expm1", "aten::special_expm1")) -def aten_special_expm1(self: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_special_expm1(self: TFloat) -> TFloat: """special_expm1(Tensor self) -> Tensor""" return op.Sub(op.Exp(self), 1) @@ -216,9 +216,7 @@ def aten_special_log_ndtr(self: TensorType) -> TensorType: @torch_op(("aten::log_softmax.int", "aten::special_log_softmax"), trace_only=True) -def aten_special_log_softmax( - self: TFloatOrBFloat16, dim: int, dtype: int = -1 -) -> TFloatOrBFloat16: +def aten_special_log_softmax(self: TFloat, dim: int, dtype: int = -1) -> TFloat: """special_log_softmax(Tensor self, int dim, *, ScalarType? dtype=None) -> Tensor""" self_is_scalar = IsScalar(self) @@ -366,7 +364,7 @@ def aten_special_xlog1py(self: TensorType, other: TensorType) -> TensorType: @torch_op(("aten::xlogy.Tensor", "aten::xlogy.Scalar_Self", "aten::xlogy.Scalar_Other")) -def aten_special_xlogy(self: TFloatOrBFloat16, other: TFloatOrBFloat16) -> TFloatOrBFloat16: +def aten_special_xlogy(self: TFloat, other: TFloat) -> TFloat: """special_xlogy(Tensor self, Tensor other) -> Tensor""" # https://pytorch.org/docs/stable/special.html#torch.special.xlogy diff --git a/onnxscript/function_libs/torch_lib/tensor_typing.py b/onnxscript/function_libs/torch_lib/tensor_typing.py index 7b5287f41..1f27c0cff 100644 --- a/onnxscript/function_libs/torch_lib/tensor_typing.py +++ b/onnxscript/function_libs/torch_lib/tensor_typing.py @@ -42,7 +42,7 @@ INT64, UINT8, ] -_FloatType = Union[FLOAT16, FLOAT, DOUBLE] +_FloatType = Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16] IntType = Union[INT8, INT16, INT32, INT64] RealType = Union[ BFLOAT16, @@ -61,7 +61,6 @@ TTensor2 = TypeVar("TTensor2", bound=_TensorType) TTensorOrString = TypeVar("TTensorOrString", bound=Union[_TensorType, STRING]) TFloat = TypeVar("TFloat", bound=_FloatType) -TFloatOrBFloat16 = TypeVar("TFloatOrBFloat16", bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16]) TFloatOrUInt8 = TypeVar("TFloatOrUInt8", bound=Union[FLOAT, FLOAT16, DOUBLE, INT8, UINT8]) TInt = TypeVar("TInt", bound=IntType) TReal = TypeVar("TReal", bound=RealType) From 1426e9f11b7bbbd9cf165d96c4c6ed9205f740d6 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 8 Oct 2024 11:56:19 -0700 Subject: [PATCH 3/6] Bump onnx-weekly in CI (#1895) To 1.18.0.dev20240930 because the previous weekly was cleaned up in pypi --- requirements/ci/requirements-onnx-weekly.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/ci/requirements-onnx-weekly.txt b/requirements/ci/requirements-onnx-weekly.txt index 2ebee9809..ea926d355 100644 --- a/requirements/ci/requirements-onnx-weekly.txt +++ b/requirements/ci/requirements-onnx-weekly.txt @@ -1 +1 @@ -onnx-weekly==1.17.0.dev20240715 +onnx-weekly==1.18.0.dev20240930 From 37b11fcff94681dea368ebfe0c4768844f2d3149 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 8 Oct 2024 14:47:36 -0700 Subject: [PATCH 4/6] [API] Create stable APIs for PyTorch 2.6 (#1896) - optimize is turned on. It will be controlled by an option in PyTorch - Remove the `_TORCH_ONNX_SAVE_EXTERNAL_DATA_WITH_IR` flag Co-authored-by: Ti-Tai Wang --- onnxscript/_framework_apis/torch_2_5.py | 72 +++++++++---------------- onnxscript/_framework_apis/torch_2_6.py | 26 +++++++++ 2 files changed, 52 insertions(+), 46 deletions(-) create mode 100644 onnxscript/_framework_apis/torch_2_6.py diff --git a/onnxscript/_framework_apis/torch_2_5.py b/onnxscript/_framework_apis/torch_2_5.py index 642660a43..eeebbb63d 100644 --- a/onnxscript/_framework_apis/torch_2_5.py +++ b/onnxscript/_framework_apis/torch_2_5.py @@ -17,17 +17,10 @@ import pathlib from typing import Callable -import onnx - from onnxscript import ir, optimizer from onnxscript.function_libs.torch_lib import registration from onnxscript.ir import _external_data -# Internal flag. Will go away. -_TORCH_ONNX_SAVE_EXTERNAL_DATA_WITH_IR = ( - os.getenv("TORCH_ONNX_OFFLOAD_EXTERNAL_DATA_WITH_IR") != "0" -) - @dataclasses.dataclass(frozen=True) class _OnnxFunctionMeta: @@ -83,45 +76,32 @@ def save_model_with_external_data(model: ir.Model, model_path: str | os.PathLike """Save the model with external data. The model is unchanged after saving.""" # TODO(#1835): Decide if we want to externalize large attributes as well - if _TORCH_ONNX_SAVE_EXTERNAL_DATA_WITH_IR: - initializer_values = tuple(model.graph.initializers.values()) - tensors = [v.const_value for v in initializer_values] - for tensor in tensors: - if tensor is None: - raise ValueError( - "The model contains uninitialized initializer values. " - "Please make sure all initializer values are initialized." - ) - destination_path = pathlib.Path(model_path) - base_dir = destination_path.parent - data_path = f"{destination_path.name}.data" - - external_tensors = _external_data.convert_tensors_to_external( - tensors, # type: ignore[arg-type] - base_dir, - data_path, - ) - - # Replace the initializer values with external tensors and save the model - for initializer, external_tensor in zip(initializer_values, external_tensors): - initializer.const_value = external_tensor - ir.save(model, model_path) - - # Restore the original initializer values so the model is unchanged - for initializer, tensor in zip(initializer_values, tensors): - initializer.const_value = tensor - - else: - destination_path = pathlib.Path(model_path) - # Create the directory if it does not exist - data_path = f"{destination_path.name}.data" - proto = ir.serde.serialize_model(model) - onnx.save_model( - proto, - model_path, - save_as_external_data=True, - location=data_path, - ) + initializer_values = tuple(model.graph.initializers.values()) + tensors = [v.const_value for v in initializer_values] + for tensor in tensors: + if tensor is None: + raise ValueError( + "The model contains uninitialized initializer values. " + "Please make sure all initializer values are initialized." + ) + destination_path = pathlib.Path(model_path) + base_dir = destination_path.parent + data_path = f"{destination_path.name}.data" + + external_tensors = _external_data.convert_tensors_to_external( + tensors, # type: ignore[arg-type] + base_dir, + data_path, + ) + + # Replace the initializer values with external tensors and save the model + for initializer, external_tensor in zip(initializer_values, external_tensors): + initializer.const_value = external_tensor + ir.save(model, model_path) + + # Restore the original initializer values so the model is unchanged + for initializer, tensor in zip(initializer_values, tensors): + initializer.const_value = tensor def get_torchlib_ops() -> list[_OnnxFunctionMeta]: diff --git a/onnxscript/_framework_apis/torch_2_6.py b/onnxscript/_framework_apis/torch_2_6.py new file mode 100644 index 000000000..ec929a1d8 --- /dev/null +++ b/onnxscript/_framework_apis/torch_2_6.py @@ -0,0 +1,26 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Stable APIs for PyTorch 2.6.""" + +from __future__ import annotations + +__all__ = [ + "check_model", + "convert_version", + "get_torchlib_ops", + "optimize", + "save_model_with_external_data", +] +from onnxscript import ir, optimizer +from onnxscript._framework_apis.torch_2_5 import ( + check_model, + convert_version, + get_torchlib_ops, + save_model_with_external_data, +) + + +def optimize(model: ir.Model) -> ir.Model: + """Optimize the model.""" + optimizer.optimize_ir(model) + return model From a7c797d7b73cde9ad964175b67dea5c7993ebddc Mon Sep 17 00:00:00 2001 From: xuzhenqi <787405797@qq.com> Date: Thu, 10 Oct 2024 04:12:49 +0800 Subject: [PATCH 5/6] Check inputs num for node matching (#1885) Fix inputs num mismatch for node matching. --------- Signed-off-by: xuzhenqi Co-authored-by: xuzhenqi Co-authored-by: Justin Chu --- onnxscript/rewriter/pattern.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/onnxscript/rewriter/pattern.py b/onnxscript/rewriter/pattern.py index be265963c..1f00840d4 100644 --- a/onnxscript/rewriter/pattern.py +++ b/onnxscript/rewriter/pattern.py @@ -959,6 +959,12 @@ def _match_node(self, pattern_node: NodePattern, node: ir.Node) -> bool: self._matched[pattern_node] = node + # TODO: Revisit this to handle optional trailing inputs better. + if len(node.inputs) != len(pattern_node.inputs): + return self.fail( + "Input nums mismatch. {len(node.inputs)} vs {len(pattern_node.inputs)}" + ) + for arg_value, arg_pattern in zip(node.inputs, pattern_node.inputs): # arg_pattern could be a Var, if it's the original arg. if arg_pattern is None: From 12f920925cb81053a4d44ab00f7a099c4f7371f3 Mon Sep 17 00:00:00 2001 From: Yuan Yao <99693700+yuanyao-nv@users.noreply.github.com> Date: Wed, 9 Oct 2024 19:26:24 -0700 Subject: [PATCH 6/6] [torchlib] Fix wrong bias shape of ConvTranspose (#1901) Previously, the bias shape of ConvTranpose was wrong since, unlike Conv, it should using the 1st dimension of weight shape and not the 0th. See described in https://github.com/microsoft/onnxscript/issues/1299 In other words, ``` if bias is None: weight_dim_0 = op.Shape(weight, start=0, end=1) bias_shape = op.Expand(weight_dim_0, op.Constant(value_ints=[1])) zero = op.CastLike(0.0, input) bias = op.Expand(zero, bias_shape) ``` should be changed to something like: ``` weight_dim_0 = op.Shape(weight, start=1, end=2) if transposed else op.Shape(weight, start=0, end=1) ``` However, I think it's more efficient to just eliminate bias altogether if it's not provided instead of filling it with zeros, since the ONNX spec allows bias to be absent. Signed-off-by: Yuan Yao --- onnxscript/function_libs/torch_lib/ops/core.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 1fc73a220..395f1fcac 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -2030,12 +2030,6 @@ def aten_convolution( stride = (stride, stride) strides = list(stride) - if bias is None: - weight_dim_0 = op.Shape(weight, start=0, end=1) - bias_shape = op.Expand(weight_dim_0, op.Constant(value_ints=[1])) - zero = op.CastLike(0.0, input) - bias = op.Expand(zero, bias_shape) - result = _aten_convolution_onnx( input, weight,