feat(ateblib):Add Ops (all, isfinite, split_with_sizes) (#513)

Co-authored-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: xiaowuhu

onnxscript/function_libs/torch_aten/ops/core.py

@@ -176,10 +176,35 @@ def aten_align_to(self: TensorType, names: Sequence[str]) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_all(self: TensorType) -> TensorType:
+@torch_op("aten::all")
+def aten_all(self: TTensor) -> BOOL:
     """all(Tensor self) -> Tensor"""
 
-    raise NotImplementedError()
+    if op.Size(op.Shape(self)) == 0:
+        result = op.Cast(self, to=BOOL.dtype)
+    else:
+        self_bool = op.Cast(self, to=BOOL.dtype)
+        self_int = op.Cast(self_bool, to=INT64.dtype)
+        result_int = op.ReduceMin(self_int, keepdims=0)
+        result = op.Cast(result_int, to=BOOL.dtype)
+
+    return result
+
+
+@torch_op("aten::all", overload=True)
+def aten_all_dim(self: TTensor, dim: int, keepdim: bool = False) -> BOOL:
+    """all(Tensor self) -> Tensor"""
+
+    if op.Size(op.Shape(self)) == 0:
+        result = op.Cast(self, to=BOOL.dtype)
+    else:
+        self_bool = op.Cast(self, to=BOOL.dtype)
+        self_int = op.Cast(self_bool, to=INT64.dtype)
+        dims = op.Reshape(dim, op.Constant(value_ints=[-1]))
+        result_int = op.ReduceMin(self_int, dims, keepdims=keepdim)
+        result = op.Cast(result_int, to=BOOL.dtype)
+
+    return result
 
 
 def aten_allclose(
@@ -2899,10 +2924,13 @@ def aten_isclose(
     raise NotImplementedError()
 
 
+@torch_op("aten::isfinite")
 def aten_isfinite(self: TensorType) -> TensorType:
     """isfinite(Tensor self) -> Tensor"""
 
-    raise NotImplementedError()
+    not_inf = op.Not(op.IsInf(self))
+    not_nan = op.Not(op.IsNaN(self))  # TODO: The test case doesnt cover this condition
+    return op.And(not_inf, not_nan)
 
 
 @torch_op("aten::isinf")
@@ -5187,10 +5215,11 @@ def aten_split_copy(self: TensorType, split_size: INT64, dim: int = 0) -> Tensor
     raise NotImplementedError()
 
 
-def aten_split_with_sizes(self: TensorType, split_sizes: INT64, dim: int = 0) -> TensorType:
+@torch_op("aten::split_with_sizes")
+def aten_split_with_sizes(self: TTensor, split_sizes: INT64, dim: int = 0) -> TTensor:
     """split_with_sizes(Tensor(a -> *) self, SymInt[] split_sizes, int dim=0) -> Tensor(a)[]"""
 
-    raise NotImplementedError()
+    return op.SplitToSequence(self, split_sizes, axis=dim)
 
 
 def aten_split_with_sizes_copy(

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -283,6 +283,8 @@ def _where_input_wrangler(
         Callable[[list[Any], dict[str, Any]], tuple[list[Any], dict[str, Any]]],
     ],
 ] = {
+    "all_dim": core_ops.aten_all_dim,
+    "all": core_ops.aten_all,
     "abs": core_ops.aten_abs,
     "acos": core_ops.aten_acos,
     "acosh": core_ops.aten_acosh,
@@ -325,6 +327,7 @@ def _where_input_wrangler(
     "full_like": core_ops.aten_full_like,
     "ge": core_ops.aten_ge,
     "gt": core_ops.aten_gt,
+    "isfinite": core_ops.aten_isfinite,
     "isinf": core_ops.aten_isinf,
     "log": core_ops.aten_log,
     "le": core_ops.aten_le,
@@ -385,6 +388,7 @@ def _where_input_wrangler(
     "sin": core_ops.aten_sin,
     "sinh": core_ops.aten_sinh,
     "softmax": special_ops.aten_special_softmax,
+    "split_with_sizes": core_ops.aten_split_with_sizes,
     "split": core_ops.aten_split,
     "sqrt": core_ops.aten_sqrt,
     "stack": core_ops.aten_stack,
@@ -478,6 +482,16 @@ def _where_input_wrangler(
 
 
 SKIP_SUBTESTS: tuple[DecorateMeta, ...] = (
+    skip(
+        "all",
+        matcher=lambda sample: not (len(sample.kwargs) == 0),
+        reason="this Aten overload only support one tensor as input by design",
+    ),
+    skip(
+        "all_dim",
+        matcher=lambda sample: not (len(sample.kwargs) > 0),
+        reason="this Aten overload only support one tensor as input and {dim,keepdim} as kwargs by design",
+    ),
     skip(
         "arange",
         matcher=lambda sample: len(sample.args) != 0,
@@ -594,6 +608,8 @@ def _where_input_wrangler(
     ),
 )
 
+duplicate_opinfo(OPS_DB, "all", ("all_dim",))
+
 duplicate_opinfo(
     OPS_DB,
     "arange",