feat(atenlib): add ops(cross_entropy_loss)

onnxscript/function_libs/torch_aten/ops/nn.py

@@ -239,17 +239,45 @@ def aten_conv_depthwise3d(
     raise NotImplementedError()
 
 
+@torch_op("aten::cross_entropy_loss", trace_only=True)
 def aten_cross_entropy_loss(
-    self: TensorType,
-    target: TensorType,
-    weight: Optional[TensorType] = None,
-    reduction: int = 1,
-    ignore_index: INT64 = -100,
-    label_smoothing: float = 0.0,
-) -> TensorType:
+    self: TFloatOrBFloat16,
+    target: Sequence[int],
+    weight: Optional[TFloatOrBFloat16] = None,
+    reduction: int = 1,  # default is 'mean'
+    ignore_index: int = -100,
+    label_smoothing: float = 0.0,  # this was ignored due to ONNX not support
+) -> TFloatOrBFloat16:
     """cross_entropy_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100, float label_smoothing=0.0) -> Tensor"""
 
-    raise NotImplementedError()
+    if reduction == 0:  # "none"
+        result = _aten_cross_entropy_loss_onnx(self, target, weight, "none", ignore_index)
+    elif reduction == 1:  # "mean"
+        result = _aten_cross_entropy_loss_onnx(self, target, weight, "mean", ignore_index)
+    else:  # "sum"
+        result = _aten_cross_entropy_loss_onnx(self, target, weight, "sum", ignore_index)
+
+    return result
+
+
+@torch_op("aten::cross_entropy_loss", overload=True)
+def _aten_cross_entropy_loss_onnx(
+    self: TFloatOrBFloat16,
+    target: Sequence[int],
+    weight: Optional[TFloatOrBFloat16],
+    reduction_str: str,
+    ignore_index: int,
+):
+    if op.OptionalHasElement(weight):
+        result, _ = op.SoftmaxCrossEntropyLoss(
+            self, target, weight, reduction=reduction_str, ignore_index=ignore_index
+        )
+    else:
+        result, _ = op.SoftmaxCrossEntropyLoss(
+            self, target, reduction=reduction_str, ignore_index=ignore_index
+        )
+
+    return result
 
 
 @torch_op("aten::elu")

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -192,6 +192,17 @@ def _cat_input_wrangler(
     return args, kwargs
 
 
+def _cross_entropy_input_wrangler(
+    args: list[Any], kwargs: dict[str, Any]
+) -> tuple[list[Any], dict[str, Any]]:
+    if "reduction" in kwargs:
+        reduction_vals = ["none", "mean", "sum"]
+        value = kwargs["reduction"]
+        idx = reduction_vals.index(value)
+        kwargs["reduction"] = idx
+    return args, kwargs
+
+
 def _dropout_input_wrangler(
     args: list[Any], kwargs: dict[str, Any]
 ) -> tuple[list[Any], dict[str, Any]]:
@@ -408,6 +419,11 @@ def _where_input_wrangler(
     "nn.functional.conv1d": core_ops.aten_conv1d,
     "nn.functional.conv2d": core_ops.aten_conv2d,
     "nn.functional.conv3d": core_ops.aten_conv3d,
+    # use cross_entropy as test case instead of cross_entropy_loss (not in OPS_DB)
+    "nn.functional.cross_entropy": (
+        nn_ops.aten_cross_entropy_loss,
+        _cross_entropy_input_wrangler,
+    ),
     "nn.functional.gelu": nn_ops.aten_gelu,
     "nn.functional.linear": nn_ops.aten_linear,
     "nn.functional.upsample_nearest2d": (
@@ -534,6 +550,11 @@ def _where_input_wrangler(
         matcher=lambda sample: isinstance(sample.kwargs.get("padding"), str),
         reason="String padding is not accepted by aten::conv2d",
     ),
+    skip(
+        "nn.functional.cross_entropy",
+        matcher=lambda sample: not isinstance(sample.kwargs.get("weight"), int),
+        reason="ONNX SoftmaxCrossEntropyLoss op only accept argument[weight] is int type",
+    ),
     skip(
         "nn.functional.dropout",
         matcher=lambda sample: len(sample.kwargs) == 0 or sample.kwargs.get("p", 0.0) > 0.0,