feat(atenlib): add ops(normal, narrow) (#440)

Author: xiaowuhu

onnxscript/function_libs/torch_aten/ops/core.py

@@ -3765,10 +3765,25 @@ def aten_nansum(
     raise NotImplementedError()
 
 
-def aten_narrow(self: TensorType, dim: int, start: INT64, length: INT64) -> TensorType:
+@torch_op("aten::narrow")
+def aten_narrow(self: TTensor, dim: INT64, start: INT64, length: INT64) -> TTensor:
     # narrow(Tensor(a) self, int dim, SymInt start, SymInt length) -> Tensor(a)
 
-    raise NotImplementedError()
+    dim_rank = op.Size(op.Shape(dim))
+    if dim_rank == 0:
+        dim = op.Reshape(dim, op.Constant(value_ints=[-1]))
+
+    start_rank = op.Size(op.Shape(start))
+    if start_rank == 0:
+        start = op.Reshape(start, op.Constant(value_ints=[-1]))
+
+    length_rank = op.Size(op.Shape(length))
+    if length_rank == 0:
+        length = op.Reshape(length, op.Constant(value_ints=[-1]))
+
+    end = op.Add(start, length)
+    result = op.Slice(self, start, end, dim)
+    return result
 
 
 def aten_narrow_copy(self: TensorType, dim: int, start: INT64, length: INT64) -> TensorType:
@@ -4036,12 +4051,20 @@ def aten_norm_except_dim(v: TensorType, pow: int = 2, dim: int = 0) -> TensorTyp
     raise NotImplementedError()
 
 
+@torch_op("aten::normal")
 def aten_normal(
-    self: TensorType, mean: float = 0.0, std: float = 1.0, generator: Optional[str] = None
-) -> TensorType:
+    self: TTensor,
+    mean: float = 0.0,
+    std: float = 1.0,
+) -> TFloat:  # type: ignore[type-var]
     # normal_functional(Tensor self, float mean=0, float std=1, *, Generator? generator=None) -> Tensor
 
-    raise NotImplementedError()
+    self_rank = op.Size(op.Shape(self))
+    if self_rank == 0:
+        self = op.Reshape(self, op.Constant(value_ints=[-1]))
+
+    result = op.RandomNormalLike(self, mean=mean, scale=std)
+    return result
 
 
 def aten_not_equal(self: TensorType, other: TensorType) -> TensorType:

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -303,6 +303,7 @@ def _where_input_wrangler(
     "minimum": core_ops.aten_minimum,
     "mm": core_ops.aten_mm,
     "mul": core_ops.aten_mul,
+    "narrow": core_ops.aten_narrow,
     # "native_dropout": core_ops.aten_native_dropout,  # native_dropout is not in OPS_DB
     "ne": core_ops.aten_ne,
     "neg": core_ops.aten_neg,
@@ -325,6 +326,7 @@ def _where_input_wrangler(
         _upsample_input_wrangler,
     ),
     "nonzero": core_ops.aten_nonzero,
+    "normal": core_ops.aten_normal,
     "ones": core_ops.aten_ones,
     "permute": core_ops.aten_permute,
     "pow": core_ops.aten_pow,
@@ -408,6 +410,8 @@ def _where_input_wrangler(
         "nn.functional.upsample_nearest2d",
         reason="enable when ONNX Runtime does support opset18",
     ),
+    xfail("normal", reason="Random numbers are not close"),
+    xfail("normal", variant_name="number_mean", reason="Random numbers are not close"),
     xfail("round", variant_name="decimals_0", reason="The op does not support decimals"),
     xfail("round", variant_name="decimals_3", reason="The op does not support decimals"),
     xfail("round", variant_name="decimals_neg_3", reason="The op does not support decimals"),
@@ -455,6 +459,11 @@ def _where_input_wrangler(
         matcher=lambda sample: sample.kwargs.get("as_tuple") is not None,
         reason="as_tuple=True is not supported",
     ),
+    skip(
+        "normal",
+        matcher=lambda sample: len(sample.args) > 0 and not isinstance(sample.args[0], float),
+        reason="ORT only accept float type for args[0] 'mean'",
+    ),
     skip(
         "nn.functional.adaptive_avg_pool1d",
         # Shape should be [N, C, D1]