feat(atenlib):add ops(isnan,issamesize,isneginf,isposinf,isnonzero)

onnxscript/function_libs/torch_aten/ops/core.py

@@ -17,6 +17,7 @@
 from onnxscript.function_libs.torch_aten.registration import torch_op
 from onnxscript.function_libs.torch_aten.tensor_typing import (
     IntType,
+    RealType,
     TFloat,
     TFloatOrBFloat16,
     TInt,
@@ -2909,10 +2910,16 @@ def aten_is_neg(self: TensorType) -> bool:
     raise NotImplementedError()
 
 
-def aten_is_nonzero(self: TensorType) -> bool:
+@torch_op("aten::is_nonzero")
+def aten_is_nonzero(self: Union[RealType, BOOL]) -> BOOL:
     """is_nonzero(Tensor self) -> bool"""
 
-    raise NotImplementedError()
+    # if size != 1, return False
+    # else [0],[True],[0.0] return True, others return False
+    result = op.Not(op.Size(self) != 1)
+    if result:
+        result = op.Cast(self, to=BOOL.dtype)
+    return result
 
 
 def aten_is_pinned(self: TensorType, device: Optional[str] = None) -> bool:
@@ -2921,10 +2928,23 @@ def aten_is_pinned(self: TensorType, device: Optional[str] = None) -> bool:
     raise NotImplementedError()
 
 
-def aten_is_same_size(self: TensorType, other: TensorType) -> bool:
+@torch_op("aten::is_same_size")
+def aten_is_same_size(self: TTensor, other: TTensor) -> BOOL:
     """is_same_size(Tensor self, Tensor other) -> bool"""
 
-    raise NotImplementedError()
+    # Cannot compare different shape of two tensors using op.Equal()
+    # So we need to compare the rank first, if rank is same, then compare shape
+    self_rank = op.Size(op.Shape(self))
+    other_rank = op.Size(op.Shape(other))
+    result = op.Equal(self_rank, other_rank)
+    if result:  # Same rank, then compare shape
+        self_shape = op.Shape(self)
+        other_shape = op.Shape(other)
+        result_bool = op.Equal(self_shape, other_shape)
+        result_int = op.Cast(result_bool, to=INT8.dtype)
+        result = op.Cast(op.ReduceMin(result_int, keepdims=0), to=BOOL.dtype)
+
+    return result
 
 
 def aten_is_set_to(self: TensorType, tensor: TensorType) -> bool:
@@ -2964,7 +2984,7 @@ def aten_isclose(
 
 
 @torch_op("aten::isfinite")
-def aten_isfinite(self: TensorType) -> TensorType:
+def aten_isfinite(self: TFloatOrBFloat16) -> BOOL:
     """isfinite(Tensor self) -> Tensor"""
 
     not_inf = op.Not(op.IsInf(self))
@@ -2979,22 +2999,25 @@ def aten_isinf(self: Union[FLOAT, DOUBLE]) -> BOOL:
     return op.IsInf(self)
 
 
-def aten_isnan(self: TensorType) -> TensorType:
+@torch_op("aten::isnan")
+def aten_isnan(self: TFloatOrBFloat16) -> BOOL:
     """isnan(Tensor self) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.IsNaN(self)
 
 
-def aten_isneginf(self: TensorType) -> TensorType:
+@torch_op("aten::isneginf")
+def aten_isneginf(self: TReal) -> BOOL:
     """isneginf(Tensor self) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.And(op.Less(self, 0), op.IsInf(self))
 
 
-def aten_isposinf(self: TensorType) -> TensorType:
+@torch_op("aten::isposinf")
+def aten_isposinf(self: TReal) -> BOOL:
     """isposinf(Tensor self) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.And(op.Greater(self, 0), op.IsInf(self))
 
 
 def aten_isreal(self: TensorType) -> TensorType:

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -404,11 +404,16 @@ def _where_input_wrangler(
     "full_like": core_ops.aten_full_like,
     "ge": core_ops.aten_ge,
     "gt": core_ops.aten_gt,
+    # "is_same_size": core_ops.aten_is_same_size,  # no test case in OPS_DB
+    # "is_nonzero": core_ops.aten_is_nonzero,  # no test case in OPS_DB
     "index_put_bool": core_ops.aten_index_put_bool,
     "index_put": core_ops.aten_index_put,
     "isclose": core_ops.aten_isclose,
     "isfinite": core_ops.aten_isfinite,
     "isinf": core_ops.aten_isinf,
+    "isnan": core_ops.aten_isnan,
+    "isneginf": core_ops.aten_isneginf,
+    "isposinf": core_ops.aten_isposinf,
     "log": core_ops.aten_log,
     "le": core_ops.aten_le,
     "log10": core_ops.aten_log10,