Update base for Update on "Auto generate OpSchema for functions | feat(op_schema)"

This change adds the capability to auto generate `OpSchema`.

### Changes

- Implement the `opschema` property in `OnnxFunction`
- Test on all torch_lib functions

### Next PR

Support trace_only functions

## Example

```python
from onnxscript.function_libs.torch_aten.ops import core, nn


print("core.aten_abs.opschema: ", core.aten_abs.opschema)

print("nn.aten_cross_entropy_loss.opschema: ", nn.aten_cross_entropy_loss.opschema)
```

Results

```
core.aten_abs.opschema:  OpSchema(
    name='aten_abs',
    domain='onnxscript.atenlib',
    since_version=1,
    doc='abs(Tensor self) -> Tensor',
    type_constraints=[OpSchema.TypeConstraintParam(type_param_str='TReal', allowed_type_strs=['tensor(float)', 'tensor(int8)', 'tensor(int16)', 'tensor(int32)', 'tensor(int64)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description='')],
    inputs=[OpSchema.FormalParameter(name='self', type_str='TReal', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    outputs=[OpSchema.FormalParameter(name='return_val', type_str='TReal', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    attributes={}
)
nn.aten_cross_entropy_loss.opschema:  OpSchema(
    name='aten_cross_entropy_loss',
    domain='onnxscript.atenlib',
    since_version=1,
    doc='cross_entropy_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100, float label_smoothing=0.0) -> Tensor',
    type_constraints=[OpSchema.TypeConstraintParam(type_param_str='TFloatOrBFloat16', allowed_type_strs=['tensor(float)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description=''), OpSchema.TypeConstraintParam(type_param_str='T1', allowed_type_strs=['tensor(float)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description='')],
    inputs=[OpSchema.FormalParameter(name='self', type_str='TFloatOrBFloat16', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>), OpSchema.FormalParameter(name='weight', type_str='T1', description='', param_option=<FormalParameterOption.Optional: 1>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    outputs=[OpSchema.FormalParameter(name='result_10', type_str='TFloatOrBFloat16', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    attributes={'ignore_index': OpSchema.Attribute(name='ignore_index', type=<AttrType.INT: 2>, description='', default_value=name: "ignore_index"
i: -100
type: INT
, required=False), 'label_smoothing': OpSchema.Attribute(name='label_smoothing', type=<AttrType.FLOAT: 1>, description='', default_value=name: "label_smoothing"
f: 0.0
type: FLOAT
, required=False), 'reduction': OpSchema.Attribute(name='reduction', type=<AttrType.INT: 2>, description='', default_value=name: "reduction"
i: 1
type: INT
, required=False), 'target': OpSchema.Attribute(name='target', type=<AttrType.INTS: 7>, description='', default_value=, required=True)}
)
```

Fixes #476

[ghstack-poisoned]

Author: justinchuby

.github/workflows/release.yaml

@@ -67,7 +67,7 @@ jobs:
           path: dist
       - name: Install wheel
         run: |
-          python -m pip install dist/*.whl
+          python -m pip install dist/*.whl --no-deps
       - name: Run tests
         run: |
           python -m pytest -v -n auto

.lintrunner.toml

@@ -1,4 +1,5 @@
 # Configuration for lintrunner https://github.com/suo/lintrunner
+merge_base_with = 'main'
 
 [[linter]]
 code = 'RUFF'

onnxscript/backend/onnx_export.py

@@ -356,12 +356,7 @@ def _python_make_node(self, onnx_node, opsets, indent=0):
             if node.op_type == "Scan":
                 return self._python_make_node_scan(node, opsets, indent=indent)
             raise RuntimeError(f"Unable to export node type {node.op_type!r} into python.")
-        if any(
-            map(
-                lambda att: hasattr(att, "g") and att.g and att.g.ByteSize() > 0,
-                node.attribute,
-            )
-        ):
+        if any(hasattr(att, "g") and att.g and att.g.ByteSize() > 0 for att in node.attribute):
             raise RuntimeError(f"Unable to export node type {node.op_type!r} into python.")
         ops = {
             "Add": "+",
@@ -438,7 +433,7 @@ def export_template(
     if hasattr(model_onnx, "functions"):
         for f in model_onnx.functions:
             unique_function_domain_version.add((f.domain, 1))
-    unique_function_domain_version_sorted = list(sorted(unique_function_domain_version))
+    unique_function_domain_version_sorted = sorted(unique_function_domain_version)
 
     if rename:
         variable_names: dict[str, str] = {}
@@ -486,7 +481,7 @@ def rename_variable(name):
             ts = _translate_type(t.type)
             its = ts.split("[", maxsplit=1)[0]
             unique_types.add(its)
-    context["unique_types"] = list(sorted(unique_types))
+    context["unique_types"] = sorted(unique_types)
 
     # functions
     functions = []

onnxscript/converter_test.py

@@ -82,7 +82,7 @@ def validate_save(
                         ort.InferenceSession(model.SerializeToString())
                     except (Fail, InvalidGraph, InvalidArgument) as e:
                         raise AssertionError(
-                            f"onnxruntime cannot load function " f"{f.name}\n--\n{model}"
+                            f"onnxruntime cannot load function {f.name}\n--\n{model}"
                         ) from e
                 if shape_inference:
                     model = onnx.shape_inference.infer_shapes(model)
@@ -423,7 +423,7 @@ def check_function(x, name, expected, eager=True):
                     y = session.run(None, {"A": x})[0]
                 except Exception as e:
                     raise AssertionError(
-                        f"Unable to run ONNX for function {name!r} " f"due to {e!r}\n{onx}."
+                        f"Unable to run ONNX for function {name!r} due to {e!r}\n{onx}."
                     ) from e
                 self.assertEqual(y.tolist(), expected)
                 f = getattr(getitem, name)
@@ -477,7 +477,7 @@ def check_function(x, name, expected, eager=True):
                     y = session.run(None, {"A": x})[0]
                 except Exception as e:
                     raise AssertionError(
-                        f"Unable to run ONNX for function {name!r} " f"due to {e!r}\n{onx}."
+                        f"Unable to run ONNX for function {name!r} due to {e!r}\n{onx}."
                     ) from e
                 self.assertEqual(y.tolist(), expected)
                 f = getattr(getitem39, name)
@@ -528,7 +528,7 @@ def check_run(self, onnxfn, inputs, expected_output):
         model = onnxfn.to_model_proto()
         session = ort.InferenceSession(model.SerializeToString())
         input_names = [x.name for x in model.graph.input]
-        input_dict = {x: value for (x, value) in zip(input_names, inputs)}
+        input_dict = dict(zip(input_names, inputs))
         output = session.run(None, input_dict)[0]
         np.testing.assert_equal(output, expected_output)
 

onnxscript/function_libs/torch_aten/ops/core.py

@@ -5365,17 +5365,58 @@ def aten_slice_copy(
     raise NotImplementedError()
 
 
+@torch_op("aten::slice_scatter", trace_only=True)
 def aten_slice_scatter(
-    self: TensorType,
-    src: TensorType,
+    self: TTensor,
+    src: TTensor,
     dim: int = 0,
     start: Optional[INT64] = None,
     end: Optional[INT64] = None,
     step: INT64 = 1,
-) -> TensorType:
+) -> TTensor:
     """slice_scatter(Tensor self, Tensor src, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor"""
 
-    raise NotImplementedError()
+    # Although 'start' and 'end' can be None in signature, but actually 'start' must be specified
+    # Assert(start is not None)
+    # And, 'end' also must be specified, and end-start must be equal to the size of 'src'
+    # Assert(end-start == shape(src) > 0)
+    # Try torch sample to get more information:
+    # https://pytorch.org/docs/master/generated/torch.slice_scatter.html?highlight=slice_scatter#torch.slice_scatter
+    # e.g. if dim=2, shape=5, permute will be [0,1]+[4]+[2,3]=[0,1,4,2,3]
+    last = len(src.shape)
+    perm = list(range(0, last))
+    perm.insert(dim, perm.pop(-1))
+    return _aten_slice_scatter_onnx(self, src, start, end, step, dim, perm)
+
+
+@torch_op("aten::slice_scatter", private=True)
+def _aten_slice_scatter_onnx(
+    self: TTensor,
+    src: TTensor,
+    start: INT64,
+    end: INT64,
+    step: INT64,
+    dim: int,
+    perm: Sequence[int],
+) -> TTensor:
+    neg_1 = op.Constant(value_ints=[-1])
+    # Get shapes expcept specifide dim
+    # e.g. if dim=2, shape=(2,3,5,7), shape_expand will be (2,3,7,1)
+    src_shape = op.Shape(src)
+    last_dim = op.Reshape(op.Size(src_shape), neg_1)
+    dim_tensor = op.Reshape(op.Constant(value_int=dim), neg_1)
+    shape_before_dim = op.Slice(src_shape, op.Constant(value_ints=[0]), dim_tensor)
+    shape_after_dim = op.Slice(src_shape, op.Add(dim_tensor, 1), last_dim)
+    shape_expand = op.Concat(
+        shape_before_dim, shape_after_dim, op.Constant(value_ints=[1]), axis=0
+    )
+    # Generate index but not finalized, need to do transpose later
+    # e.g. [[0,1,2],[0,1,2],[0,1,2]...,[0,1,2]], total count = 2x3x7
+    index_base = op.Range(start, end, step)  # e.g. [0,1,2]
+    index_expand = op.Expand(index_base, shape_expand)
+    indices = op.Transpose(index_expand, perm=perm)
+
+    return op.ScatterElements(self, indices, src, axis=dim)
 
 
 def aten_slogdet(self: TensorType) -> tuple[TensorType, TensorType]:
@@ -6043,96 +6084,10 @@ def aten_var(self: TensorType, unbiased: bool = True) -> TensorType:
     raise NotImplementedError()
 
 
-@torch_op("aten::var_mean", trace_only=True)
-def aten_var_mean(self: TReal, unbiased: bool = True) -> Tuple[TReal, TReal]:
+def aten_var_mean(self: TensorType, unbiased: bool = True) -> tuple[TensorType, TensorType]:
     """var_mean(Tensor self, bool unbiased=True) -> (Tensor, Tensor)"""
 
-    # 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_mean_onnx(self, correction=int(unbiased), keepdim=False)
-
-
-@torch_op("aten::var_mean", overload=True, trace_only=True)
-def aten_var_mean_dim(
-    self: TReal, dim: Optional[int], unbiased: bool = True, keepdim: bool = False
-) -> Tuple[TReal, TReal]:
-    """var_mean.dim(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False) -> (Tensor, Tensor)"""
-
-    # Although dim is Optional in signature, but we assume it must has value for this overload
-    # Assert(dim is not None)
-    if isinstance(dim, Tuple):
-        dim_tensor = op.Constant(value_ints=dim)
-    else:
-        dim_tensor = op.Constant(value_int=dim)
-    return _aten_var_mean_dim_onnx(self, dim_tensor, correction=int(unbiased), keepdim=keepdim)
-
-
-@torch_op("aten::var_mean", overload=True, trace_only=True)
-def aten_var_mean_correction(
-    self: TReal,
-    dim: Optional[int] = None,
-    correction: Optional[int] = None,
-    keepdim: bool = False,
-) -> Tuple[TReal, TReal]:
-    """var_mean.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> (Tensor, Tensor)"""
-
-    if correction is None:
-        correction = 1
-
-    if dim is None:
-        var, mean = _aten_var_mean_onnx(self, correction, keepdim)
-    else:
-        if isinstance(dim, Tuple):
-            dim_tensor = op.Constant(value_ints=dim)
-        else:
-            dim_tensor = op.Constant(value_int=dim)
-        var, mean = _aten_var_mean_dim_onnx(self, dim_tensor, correction, keepdim)
-    return var, mean
-
-
-@torch_op("aten::var_mean", private=True)
-def _aten_var_mean_onnx(
-    self: TReal, correction: int = 1, keepdim: bool = False
-) -> Tuple[TReal, TReal]:
-    # Compute mean and var
-    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:
-        self_shape = op.Shape(self)
-        numel_int = op.ReduceProd(self_shape, keepdims=0)
-        numel_float = op.Cast(numel_int, to=FLOAT.dtype)
-        mul = op.Mul(var, numel_float)
-        sub = op.Sub(numel_int, correction)
-        var = op.Div(mul, op.Cast(sub, to=FLOAT.dtype))
-
-    return var, mean
-
-
-@torch_op("aten::var_mean", private=True)
-def _aten_var_mean_dim_onnx(
-    self: TReal, dim: INT64, correction: int, keepdim: bool = False
-) -> Tuple[TReal, TReal]:
-    if op.Size(op.Shape(dim)) == 0:
-        dim = op.Unsqueeze(dim, axes=0)
-    # Computer mean and var
-    mean = op.ReduceMean(self, dim, keepdims=keepdim)
-    sub_mean = op.Sub(self, op.ReduceMean(self, dim, keepdims=1))
-    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:
-        self_shape = op.Shape(self)
-        dim_size = op.Gather(self_shape, dim, axis=0)
-        numel_int = op.ReduceProd(dim_size, keepdims=0)
-        numel_float = op.Cast(numel_int, to=FLOAT.dtype)
-        mul = op.Mul(var, numel_float)
-        sub = op.Sub(numel_int, correction)
-        var = op.Div(mul, op.Cast(sub, to=FLOAT.dtype))
-
-    return var, mean
+    raise NotImplementedError()
 
 
 def aten_vdot(self: TensorType, other: TensorType) -> TensorType:

onnxscript/function_libs/torch_aten/ops/nn.py

@@ -529,10 +529,11 @@ def aten_hardswish_backward(grad_output: TensorType, self: TensorType) -> Tensor
     raise NotImplementedError()
 
 
-def aten_hardtanh(self: TensorType, min_val: float = -1.0, max_val: float = 1.0) -> TensorType:
+@torch_op("aten::hardtanh")
+def aten_hardtanh(self: TReal, min_val: float = -1.0, max_val: float = 1.0) -> TReal:
     """hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.Clip(self, min_val, max_val)
 
 
 def aten_hardtanh_backward(

onnxscript/tests/eager_test.py

@@ -22,7 +22,7 @@ def _fft(x, fft_length, axis=-1):
     tr = np.transpose(merged, list(perm))
     if tr.shape[-1] != 2:
         raise AssertionError(
-            f"Unexpected shape {tr.shape}, x.shape={x.shape} " f"fft_length={fft_length}."
+            f"Unexpected shape {tr.shape}, x.shape={x.shape} fft_length={fft_length}."
         )
     return tr
 
@@ -48,7 +48,7 @@ def _ifft(x, fft_length, axis=-1):
     tr = np.transpose(merged, list(perm))
     if tr.shape[-1] != 2:
         raise AssertionError(
-            f"Unexpected shape {tr.shape}, x.shape={x.shape} " f"fft_length={fft_length}."
+            f"Unexpected shape {tr.shape}, x.shape={x.shape} fft_length={fft_length}."
         )
     return tr
 

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -610,6 +610,7 @@ def _where_input_wrangler(
     "nn.functional.dropout": (core_ops.aten_dropout, _dropout_input_wrangler),
     "nn.functional.elu": nn_ops.aten_elu,
     "nn.functional.embedding": (core_ops.aten_embedding, _embedding_input_wrangler),
+    "nn.functional.hardtanh": nn_ops.aten_hardtanh,
     "nn.functional.leaky_relu": nn_ops.aten_leaky_relu,
     "nn.functional.logsigmoid": nn_ops.aten_log_sigmoid,
     "nn.functional.nll_loss_weight": (nn_ops.aten_nll_loss_weight, _nll_loss_input_wrangler),
@@ -731,12 +732,10 @@ def _where_input_wrangler(
     ),
     "ones_like": core_ops.aten_ones_like,
     "scatter_reduce": (core_ops.aten_scatter_reduce, _scatter_reduce_input_wrangler),
+    "slice_scatter": core_ops.aten_slice_scatter,
     "slice": core_ops.aten_slice,
     "sum": (core_ops.aten_sum_dim_IntList, _sum_input_wrangler),
     "transpose": core_ops.aten_transpose,
-    "var_mean": core_ops.aten_var_mean,
-    "var_mean_dim": core_ops.aten_var_mean_dim,
-    "var_mean_correction": core_ops.aten_var_mean_correction,
     "zeros_like": core_ops.aten_zeros_like,
 }
 
@@ -1194,27 +1193,6 @@ def _where_input_wrangler(
         matcher=lambda sample: not (len(sample.args) > 0 and isinstance(sample.args[0], int)),
         reason="this Aten overload only support one tensor as input and one int as args by design",
     ),
-    skip(
-        "var_mean",
-        # kwargs is 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",
-    ),
-    skip(
-        "var_mean_dim",
-        # 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",
-    ),
-    skip(
-        "var_mean_correction",
-        # 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",
-    ),
     skip(
         "unflatten",
         matcher=lambda sample: any(dim == 0 for dim in sample.input.shape),
@@ -1284,15 +1262,6 @@ def _where_input_wrangler(
 
 duplicate_opinfo(OPS_DB, "squeeze", ("squeeze_dim",))
 
-duplicate_opinfo(
-    OPS_DB,
-    "var_mean",
-    (
-        "var_mean_dim",
-        "var_mean_correction",
-    ),
-)
-
 
 # END OF SECTION TO MODIFY #####################################################
 

opgen/onnx_opset_builder.py

@@ -123,9 +123,7 @@ def _make_opset_module(self, domain: str, version: int):
                 cg.FunctionDef(
                     "__new__",
                     cg.Arg("cls"),
-                    body=cg.ThunkStmt(
-                        f"return Opset.__new__(cls, " f"{domain!r}, {version!r})"
-                    ),
+                    body=cg.ThunkStmt(f"return Opset.__new__(cls, {domain!r}, {version!r})"),
                 ),
                 cg.FunctionDef(
                     "__init__", cg.Arg("self"), body=cg.ThunkStmt("super().__init__()")

pyproject.toml

@@ -106,19 +106,23 @@ convention = "google"
 [tool.ruff]
 target-version = "py38"
 select = [
+    "B", # flake8-bugbear
+    "C4", # flake8-comprehensions
     "D", # pydocstyle
     "E", # pycodestyle
     "F", # Pyflakes
-    "W", # pycodestyle
-    "B", # flake8-bugbear
+    "G", # flake8-logging-format
+    "ISC", # flake8-implicit-str-concat
     "N", # pep8-naming
     "NPY", # modern numpy
-    "YTT", # flake8-2020
     "RUF", # Ruff-specific rules
-    "UP", # pyupgrade
     "TID252", # Disallow relative imports
+    "UP", # pyupgrade
+    "W", # pycodestyle
+    "YTT", # flake8-2020
 ]
 ignore = [
+    "C408", # Sometimes it is preferable when we construct kwargs
     "D1", # D1 is for missing docstrings, which is not yet enforced.
     "D202", # D202 Too strict. "No blank lines allowed after function docstring"
     "D205", # D205 Too strict. "1 blank line required between summary line and description"

requirements-dev.txt

@@ -27,5 +27,5 @@ pyyaml
 torch>=1.13
 
 # Lint
-lintrunner
+lintrunner>=0.10.7
 lintrunner_adapters>=0.7.0