Merge branch 'main' into xiaowu/addOps(0216)

Author: xiaowuhu

onnxscript/function_libs/torch_aten/ops/core.py

@@ -269,10 +269,32 @@ def aten_angle(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_any(self: TensorType) -> TensorType:
+@torch_op("aten::any", trace_only=True)
+def aten_any(self: TTensor, dim: Optional[int] = None, keepdim: bool = True) -> BOOL:
     """any(Tensor self) -> Tensor"""
 
-    raise NotImplementedError()
+    minus_1 = op.Constant(value_ints=[-1])
+    self_rank = op.Size(op.Shape(self))
+    if self_rank == 0:
+        self = op.Reshape(self, minus_1)
+
+    zero = op.Constant(value_float=0.0)
+    result = op.Not(op.Equal(self, zero))
+    # because op.ReduceMax() cannot calculate BOOL value
+    result_float = op.Cast(result, to=FLOAT.dtype)
+
+    if op.OptionalHasElement(dim):
+        dim = op.Reshape(dim, minus_1)
+        dims = op.Cast(dim, to=INT64.dtype)
+        result_max = op.ReduceMax(result_float, dims, keepdims=keepdim, noop_with_empty_axes=0)
+    else:
+        result_max = op.ReduceMax(result_float, keepdims=0, noop_with_empty_axes=0)
+
+    result = op.Greater(result_max, zero)
+    if self_rank == 0:
+        result = op.Squeeze(result)
+
+    return result
 
 
 def _range_supported(dtype: int) -> bool:
@@ -2044,10 +2066,13 @@ def aten_expand(self: TTensor, size: TInt) -> TTensor:
     return op.Expand(self, size)
 
 
-def aten_expand_as(self: TensorType, other: TensorType) -> TensorType:
+@torch_op("aten::expand_as")
+def aten_expand_as(self: TTensor, other: TTensor) -> TTensor:
     """expand_as(Tensor(a) self, Tensor other) -> Tensor(a)"""
 
-    raise NotImplementedError()
+    shape = op.Shape(other)
+    result = op.Expand(self, shape)
+    return result
 
 
 def aten_expand_copy(self: TensorType, size: INT64, implicit: bool = False) -> TensorType:

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -314,6 +314,7 @@ def _where_input_wrangler(
     "exp": core_ops.aten_exp,
     "exp2": core_ops.aten_exp2,
     "expand": core_ops.aten_expand,
+    "expand_as": core_ops.aten_expand_as,
     "erf": core_ops.aten_erf,
     "fill": core_ops.aten_fill,
     "fmod": core_ops.aten_fmod,
@@ -398,6 +399,7 @@ def _where_input_wrangler(
 ] = {
     "amax": core_ops.aten_amax,
     "amin": core_ops.aten_amin,
+    "any": core_ops.aten_any,  # TODO: add more testcase which element is [0.0, 0.1, -0.1, 0.0] etc.
     "arange_start_step": core_ops.aten_arange_start_step,
     "arange_start": core_ops.aten_arange_start,
     "arange": core_ops.aten_arange,