Merge branch 'main' into xiaowu/trySome2n

Author: xiaowuhu

onnxscript/function_libs/torch_aten/ops/core.py

@@ -48,20 +48,6 @@ def aten_acosh(self: TFloat) -> TFloat:
     return op.Acosh(self)
 
 
-def aten_adaptive_avg_pool1d(self: TensorType, output_size: Sequence[int]) -> TensorType:
-    # adaptive_avg_pool1d(Tensor self, int[1] output_size) -> Tensor
-
-    raise NotImplementedError()
-
-
-def aten_adaptive_max_pool1d(
-    self: TensorType, output_size: Sequence[int]
-) -> tuple[TensorType, TensorType]:
-    # adaptive_max_pool1d(Tensor self, int[1] output_size) -> (Tensor, Tensor)
-
-    raise NotImplementedError()
-
-
 @torch_op("aten::add")
 def aten_add(self: TReal, other: TReal, alpha: float = 1) -> TReal:
     # add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
@@ -198,20 +184,20 @@ def aten_alpha_dropout(input: TensorType, p: float, train: bool) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_amax(
-    self: TensorType, dim: Optional[Sequence[int]] = None, keepdim: bool = False
-) -> TensorType:
+# @torch_op("aten::amax")  # FIXME: Uncomment when CI uses onnx 1.13
+def aten_amax(self: TReal, dim: INT64, keepdim: int = 0) -> TReal:
     # amax(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor
 
-    raise NotImplementedError()
+    # TODO(justinchuby): Make dim optional, keepdim bool
+    return op.ReduceMax(self, dim, keepdims=keepdim)
 
 
-def aten_amin(
-    self: TensorType, dim: Optional[Sequence[int]] = None, keepdim: bool = False
-) -> TensorType:
+# @torch_op("aten::amin")  # FIXME: Uncomment when CI uses onnx 1.13
+def aten_amin(self: TReal, dim: INT64, keepdim: int = 0) -> TReal:
     # amin(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor
 
-    raise NotImplementedError()
+    # TODO(justinchuby): Make dim optional, keepdim bool
+    return op.ReduceMin(self, dim, keepdims=keepdim)
 
 
 def aten_aminmax(
@@ -2190,10 +2176,21 @@ def aten_index_reduce(
     raise NotImplementedError()
 
 
-def aten_index_select(self: TensorType, dim: int, index: TensorType) -> TensorType:
+# FIXME(#277): Script when attributes can come before inputs
+@torch_op("aten::index_select", trace_only=True)
+def aten_index_select(self: TTensor, dim: int, index: TInt) -> TTensor:
     # index_select(Tensor self, int dim, Tensor index) -> Tensor
 
-    raise NotImplementedError()
+    if op.Size(op.Shape(self)) == 0:
+        result = self
+    else:
+        # Index can be a scalar. Reshape it to a rank 1 tensor.
+        index = op.Reshape(index, op.Constant(value_floats=[-1]))
+        index = op.Cast(index, to=INT64.dtype)
+
+        result = op.Gather(self, index, axis=dim)
+
+    return result
 
 
 def aten_index_select_backward(
@@ -4194,10 +4191,11 @@ def aten_rsqrt(self: TFloatOrBFloat16) -> TFloatOrBFloat16:
     return op.Reciprocal(op.Sqrt(self))
 
 
-def aten_rsub(self: TensorType, other: TensorType, alpha: float = 1) -> TensorType:
+@torch_op("aten::rsub")
+def aten_rsub(self: TReal, other: TReal, alpha: float = 1.0) -> TReal:
     # rsub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
 
-    raise NotImplementedError()
+    return op.Sub(other, op.Mul(self, alpha))
 
 
 def aten_scalar_tensor(s: float) -> TensorType:
@@ -4698,11 +4696,26 @@ def aten_trace_backward(grad: TensorType, sizes: INT64) -> TensorType:
     raise NotImplementedError()
 
 
+@torch_op("aten::transpose", trace_only=True)
 def aten_transpose(self, dim0: int, dim1: int):
     # transpose.int(Tensor(a) self, int dim0, int dim1) -> Tensor(a)
 
     # FIXME(justinchuby): onnxscript raises Unsupported expression type
-    return op.Transpose(self, [dim0, dim1])
+    # Script the function when this is fixed
+    self_rank = op.Size(op.Shape(self))
+
+    if self_rank == 0:
+        result = self
+    else:
+        # Python code, change when onnxscript supports this
+        self_rank_val = self_rank.value  # type: ignore[attr-defined]
+        dims = list(range(self_rank_val))
+        dims[dim0], dims[dim1] = dims[dim1], dims[dim0]
+        # Python code ends
+
+        result = op.Transpose(self, perm=dims)
+
+    return result
 
 
 def aten_triangular_solve(

onnxscript/function_libs/torch_aten/ops/nn.py

@@ -23,15 +23,65 @@
 from onnxscript.onnx_types import TensorType
 
 
-def aten_adaptive_avg_pool2d(self: TensorType, output_size: INT64) -> TensorType:
+@torch_op("aten::aten_adaptive_avg_pool1d")
+def aten_adaptive_avg_pool1d(self: TFloat, output_size: INT64[1]) -> TFloat:
+    # adaptive_avg_pool1d(Tensor self, int[1] output_size) -> Tensor
+
+    # assert output_size == [1]
+    # TODO(justinchuby): Specify input constraints
+
+    if op.Size(op.Shape(self)) == 2:
+        # Unbatched case
+        self = op.Unsqueeze(self, op.Constant(value_ints=[0]))
+        pooled = op.GlobalAveragePool(self)
+        result = op.Squeeze(pooled, op.Constant(value_ints=[0]))
+    else:
+        result = op.GlobalAveragePool(self)
+
+    return result
+
+
+@torch_op("aten::aten_adaptive_avg_pool2d")
+def aten_adaptive_avg_pool2d(self: TFloat, output_size: INT64[2]) -> TFloat:
     # adaptive_avg_pool2d(Tensor self, SymInt[2] output_size) -> Tensor
 
-    raise NotImplementedError()
+    # assert output_size == [1, 1]
+    # TODO(justinchuby): Specify input constraints
+
+    if op.Size(op.Shape(self)) == 3:
+        # Unbatched case
+        self = op.Unsqueeze(self, op.Constant(value_ints=[0]))
+        pooled = op.GlobalAveragePool(self)
+        result = op.Squeeze(pooled, op.Constant(value_ints=[0]))
+    else:
+        result = op.GlobalAveragePool(self)
+
+    return result
 
 
-def aten_adaptive_avg_pool3d(self: TensorType, output_size: INT64) -> TensorType:
+@torch_op("aten::aten_adaptive_avg_pool3d")
+def aten_adaptive_avg_pool3d(self: TFloat, output_size: INT64[3]) -> TFloat:
     # adaptive_avg_pool3d(Tensor self, SymInt[3] output_size) -> Tensor
 
+    # assert output_size == [1, 1, 1]
+    # TODO(justinchuby): Specify input constraints
+
+    if op.Size(op.Shape(self)) == 4:
+        # Unbatched case
+        self = op.Unsqueeze(self, op.Constant(value_ints=[0]))
+        pooled = op.GlobalAveragePool(self)
+        result = op.Squeeze(pooled, op.Constant(value_ints=[0]))
+    else:
+        result = op.GlobalAveragePool(self)
+
+    return result
+
+
+def aten_adaptive_max_pool1d(
+    self: TensorType, output_size: Sequence[int]
+) -> tuple[TensorType, TensorType]:
+    # adaptive_max_pool1d(Tensor self, int[1] output_size) -> (Tensor, Tensor)
+
     raise NotImplementedError()
 
 
@@ -1162,15 +1212,29 @@ def aten_upsample_nearest1d_backward(
     raise NotImplementedError()
 
 
+@torch_op("aten::upsample_nearest2d")
 def aten_upsample_nearest2d(
-    self: TensorType,
-    output_size: INT64,
+    self: TReal,
+    size: INT64,
     scales_h: Optional[float] = None,
     scales_w: Optional[float] = None,
-) -> TensorType:
+) -> TReal:
     # upsample_nearest2d(Tensor self, SymInt[2] output_size, float? scales_h=None, float? scales_w=None) -> Tensor
 
-    raise NotImplementedError()
+    self_shape = op.Shape(self)
+    batch_channel = self_shape[:2]  # type: ignore[index]
+    output_size = op.Concat(batch_channel, size, axis=0)
+
+    # TODO(justinchuby): Conditionally use scales
+
+    return op.Resize(
+        self,
+        None,
+        None,
+        output_size,
+        mode="nearest",
+        coordinate_transformation_mode="asymmetric",
+    )
 
 
 def aten_upsample_nearest2d_backward(

onnxscript/function_libs/torch_aten/registration.py

@@ -48,19 +48,33 @@ def __repr__(self):
 
 
 def torch_op(
-    name, overload: bool = False, registry: Optional[Registry] = None
-) -> Callable[[Callable[..., Any]], onnxscript.OnnxFunction]:
-    """Register a torch op."""
+    name,
+    *,
+    overload: bool = False,
+    registry: Optional[Registry] = None,
+    trace_only: bool = False,
+) -> Callable[[Callable[..., Any]], onnxscript.OnnxFunction | Callable[..., Any]]:
+    """Register a torch op.
+
+    Args:
+        name: ATen name of the function. E.g. "aten::add".
+        overload: Whether the function is an overload (not default).
+        registry: Registry to register the function to. If None, the default registry is used.
+        trace_only: Whether the function should only be traced and not compiled.
+    """
     if registry is None:
         registry = default_registry
 
-    def wrapper(func: Callable[..., Any]) -> onnxscript.OnnxFunction:
+    def wrapper(func: Callable[..., Any]) -> onnxscript.OnnxFunction | Callable[..., Any]:
 
-        # Compile the function
-        compiled = onnxscript.script()(func)
+        if trace_only:
+            processed_func = func
+        else:
+            # Compile the function
+            processed_func = onnxscript.script()(func)
 
         assert registry is not None
-        registry.register(compiled, name, overload=overload)
-        return compiled
+        registry.register(processed_func, name, overload=overload)
+        return processed_func
 
     return wrapper