onnxscript 0.1.0.dev20240116__tar.gz → 0.1.0.dev20240118__tar.gz

{onnxscript-0.1.0.dev20240116/onnxscript.egg-info → onnxscript-0.1.0.dev20240118}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxscript
-Version: 0.1.0.dev20240116
+Version: 0.1.0.dev20240118
 Summary: Naturally author ONNX functions and models using a subset of Python
 Home-page: https://onnxscript.ai/
 Author-email: Microsoft Corporation <onnx@microsoft.com>

{onnxscript-0.1.0.dev20240116 → onnxscript-0.1.0.dev20240118}/onnxscript/function_libs/torch_lib/ops/__init__.py

@@ -7,6 +7,7 @@ __all__ = [
     "prims",
     "sparse",
     "special",
+    "vision",
 ]
 
-from . import core, fft, linalg, nested, nn, prims, sparse, special
+from . import core, fft, linalg, nested, nn, prims, sparse, special, vision

{onnxscript-0.1.0.dev20240116 → onnxscript-0.1.0.dev20240118}/onnxscript/function_libs/torch_lib/ops/nn.py

@@ -2295,38 +2295,45 @@ def aten_upsample_bicubic2d_backward(
 @torch_op("aten::upsample_bilinear2d", trace_only=True)
 def aten_upsample_bilinear2d(
     self: TReal,
-    output_size: Optional[INT64] = None,
+    output_size: INT64,
+    align_corners: bool,
     scales_h: Optional[float] = None,
     scales_w: Optional[float] = None,
-    align_corners: bool = True,  # pylint: disable=unused-argument
 ) -> TReal:
     """upsample_bilinear2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor"""
 
+    coordinate_transformation_mode = "align_corners" if align_corners else "pytorch_half_pixel"
     if output_size is not None:
-        result = _aten_upsample_bilinear2d_output_size(self, output_size)
+        result = _aten_upsample_bilinear2d_output_size(
+            self, output_size, coordinate_transformation_mode
+        )
     else:
         assert scales_h is not None
         assert scales_h == scales_w, f"scale_h({scales_h}) != scale_w({scales_w})"
-        result = _aten_upsample_bilinear2d_scales(self, scales_h, scales_w)
+        result = _aten_upsample_bilinear2d_scales(
+            self, scales_h, scales_w, coordinate_transformation_mode
+        )
     return result
 
 
 @torch_op("aten::upsample_bilinear2d.vec", trace_only=True)
 def aten_upsample_bilinear2d_vec(
     self: TReal,
-    output_size: Optional[INT64] = None,
-    align_corners: bool = True,
-    scale_factors: Optional[Sequence[float]] = None,
+    output_size: Optional[INT64],
+    align_corners: bool,
+    scale_factors: Optional[Sequence[float]],
 ) -> TReal:
+    """upsample_bilinear2d.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor"""
     scales_h = scale_factors[0] if scale_factors is not None else None
     scales_w = scale_factors[1] if scale_factors is not None else None
-    return aten_upsample_bilinear2d(self, output_size, scales_h, scales_w, align_corners)
+    return aten_upsample_bilinear2d(self, output_size, align_corners, scales_h, scales_w)
 
 
 @torch_op("aten::upsample_bilinear2d", private=True)
 def _aten_upsample_bilinear2d_output_size(
     self: TReal,
     output_size: INT64,
+    coordinate_transformation_mode: str,
 ) -> TReal:
     """upsample_bilinear2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor"""
 
@@ -2341,7 +2348,8 @@ def _aten_upsample_bilinear2d_output_size(
         None,
         output_size,
         mode="linear",
-        coordinate_transformation_mode="align_corners",
+        coordinate_transformation_mode=coordinate_transformation_mode,
+        nearest_mode="floor",
     )
 
 
@@ -2350,6 +2358,7 @@ def _aten_upsample_bilinear2d_scales(
     self: TReal,
     scales_h: float,
     scales_w: float,
+    coordinate_transformation_mode: str,
 ) -> TReal:
     """upsample_bilinear2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor"""
 
@@ -2366,7 +2375,8 @@ def _aten_upsample_bilinear2d_scales(
         scales,  # format should be: [1.0, 1.0, scale_h, scale_w]
         None,
         mode="linear",
-        coordinate_transformation_mode="align_corners",
+        coordinate_transformation_mode=coordinate_transformation_mode,
+        nearest_mode="floor",
     )
 
 

onnxscript-0.1.0.dev20240118/onnxscript/function_libs/torch_lib/ops/vision.py

@@ -0,0 +1,24 @@
+# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+# mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value"
+"""torchvision operators."""
+from __future__ import annotations
+
+from onnxscript.function_libs.torch_lib.registration import torch_op
+from onnxscript.onnx_opset import opset18 as op
+from onnxscript.onnx_types import FLOAT, INT64
+
+_INT64_MAX = 0x7FFFFFFFFFFFFFFF
+
+
+@torch_op("torchvision::nms")
+def torchvision_nms(boxes: FLOAT, scores: FLOAT, iou_threshold: float) -> INT64:
+    # boxes: [num_batches, spatial_dimension, 4]
+    boxes = op.Unsqueeze(boxes, [0])
+    # scores: [num_batches, num_classes, spatial_dimension]
+    scores = op.Unsqueeze(scores, [0, 1])
+    # nms_out: [num_selected_indices, 3] where each column is [batch_index, class_index, box_index]
+    nms_out = op.NonMaxSuppression(boxes, scores, _INT64_MAX, iou_threshold)
+    return op.Reshape(op.Slice(nms_out, axes=[1], starts=[2], ends=[3]), [-1])

{onnxscript-0.1.0.dev20240116 → onnxscript-0.1.0.dev20240118}/onnxscript/tests/function_libs/torch_lib/extra_opinfo.py

@@ -8,6 +8,7 @@ import itertools
 from typing import Any, List
 
 import torch
+import torchvision
 from torch import testing as torch_testing
 from torch.testing._internal import (
     common_device_type,
@@ -997,6 +998,27 @@ def sample_inputs__native_batch_norm_legit_no_stats(
             )
 
 
+def sample_inputs_non_max_suppression(op_info, device, dtype, requires_grad, **kwargs):
+    del op_info
+    del kwargs
+    boxes = torch.tensor(
+        [
+            [0.0, 0.0, 10.0, 10.0],
+            [10.0, 10.0, 20.0, 20.0],
+            [32.0, 32.0, 40.0, 52.0],
+        ],
+        device=device,
+        dtype=dtype,
+        requires_grad=requires_grad,
+    )
+    scores = torch.tensor(
+        [0.8, 0.4, 0.6], device=device, dtype=dtype, requires_grad=requires_grad
+    )
+
+    for iou_threshold in (0.3, 0.5, 0.7, 0.9):
+        yield opinfo_core.SampleInput(boxes, args=(scores, iou_threshold))
+
+
 def sample_inputs_normal_tensor_float(op_info, device, dtype, requires_grad, **kwargs):
     del op_info
     del requires_grad
@@ -1961,4 +1983,11 @@ OP_DB: List[opinfo_core.OpInfo] = [
         supports_autograd=False,
         supports_out=False,
     ),
+    opinfo_core.OpInfo(
+        "torchvision.ops.nms",
+        op=torchvision.ops.nms,
+        dtypes=common_dtype.floating_types(),
+        sample_inputs_func=sample_inputs_non_max_suppression,
+        supports_out=False,
+    ),
 ]

{onnxscript-0.1.0.dev20240116 → onnxscript-0.1.0.dev20240118}/onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -52,6 +52,7 @@ from onnxscript.function_libs.torch_lib.ops import fft as fft_ops
 from onnxscript.function_libs.torch_lib.ops import linalg as linalg_ops
 from onnxscript.function_libs.torch_lib.ops import nn as nn_ops
 from onnxscript.function_libs.torch_lib.ops import special as special_ops
+from onnxscript.function_libs.torch_lib.ops import vision as vision_ops
 from onnxscript.tests.function_libs.torch_lib import extra_opinfo, ops_test_common
 
 # Create a copy of the op_db to modify
@@ -417,9 +418,21 @@ def _sum_input_wrangler(
 def _upsample_bilinear2d_input_wrangler(
     args: list[Any], kwargs: dict[str, Any]
 ) -> tuple[list[Any], dict[str, Any]]:
+    # Wrangler for the signature difference between
+    #   'nn.functional.upsample_bilinear'
+    # and
+    #   'aten::upsample_bilinear2d'
+    # https://pytorch.org/docs/stable/generated/torch.nn.functional.upsample_bilinear.html
     if "size" in kwargs:
         args.append(np.array(kwargs["size"], dtype=np.int64))
         del kwargs["size"]  # promote tensor type kwargs to args
+    else:
+        args.append(None)
+    if "align_corners" in kwargs:
+        args.append(kwargs["align_corners"])
+        del kwargs["align_corners"]
+    else:
+        args.append(True)  # Fill in the default value
     if "scale_factor" in kwargs:
         kwargs["scales_h"] = kwargs["scale_factor"]
         kwargs["scales_w"] = kwargs["scale_factor"]
@@ -430,12 +443,26 @@ def _upsample_bilinear2d_input_wrangler(
 def _upsample_bilinear2d_vec_input_wrangler(
     args: list[Any], kwargs: dict[str, Any]
 ) -> tuple[list[Any], dict[str, Any]]:
+    # Wrangler for the signature difference between
+    #   'nn.functional.upsample_bilinear'
+    # and
+    #   'aten::upsample_bilinear2d.vec'
+    # https://pytorch.org/docs/stable/generated/torch.nn.functional.upsample_bilinear.html
     if "size" in kwargs:
         args.append(np.array(kwargs["size"], dtype=np.int64))
         del kwargs["size"]  # promote tensor type kwargs to args
+    else:
+        args.append(None)
+    if "align_corners" in kwargs:
+        args.append(kwargs["align_corners"])
+        del kwargs["align_corners"]
+    else:
+        args.append(True)  # Fill in the default value
     if "scale_factor" in kwargs:
-        kwargs["scale_factors"] = [kwargs["scale_factor"]] * 2
-        del kwargs["scale_factor"]  # adapt the function signature
+        args.append([kwargs["scale_factor"]] * 2)
+        del kwargs["scale_factor"]  # promote tensor type kwargs to args
+    else:
+        args.append(None)
     return args, kwargs
 
 
@@ -2300,6 +2327,7 @@ TESTED_TORCHLIB_OPS: tuple[TorchLibOpInfo, ...] = (
         reason="this Aten overload only support when correction attribute exists",
     ),
     TorchLibOpInfo("zeros_like", core_ops.aten_zeros_like, trace_only=True),
+    TorchLibOpInfo("torchvision.ops.nms", vision_ops.torchvision_nms),
 )
 
 ops_test_common.duplicate_opinfo(OPS_DB, "all", ("all_dim", "all_dims"))

{onnxscript-0.1.0.dev20240116 → onnxscript-0.1.0.dev20240118/onnxscript.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxscript
-Version: 0.1.0.dev20240116
+Version: 0.1.0.dev20240118
 Summary: Naturally author ONNX functions and models using a subset of Python
 Home-page: https://onnxscript.ai/
 Author-email: Microsoft Corporation <onnx@microsoft.com>

{onnxscript-0.1.0.dev20240116 → onnxscript-0.1.0.dev20240118}/onnxscript.egg-info/SOURCES.txt

@@ -124,6 +124,7 @@ onnxscript/function_libs/torch_lib/ops/nn.py
 onnxscript/function_libs/torch_lib/ops/prims.py
 onnxscript/function_libs/torch_lib/ops/sparse.py
 onnxscript/function_libs/torch_lib/ops/special.py
+onnxscript/function_libs/torch_lib/ops/vision.py
 onnxscript/onnx_opset/__init__.py
 onnxscript/onnx_opset/_impl/opset1.py
 onnxscript/onnx_opset/_impl/opset10.py