feat(atenlib): add ops(cross_entropy_loss)

onnxscript/function_libs/torch_aten/ops/nn.py

@@ -239,17 +239,26 @@ def aten_conv_depthwise3d(
     raise NotImplementedError()
 
 
-def aten_cross_entropy_loss(
-    self: TensorType,
-    target: TensorType,
-    weight: Optional[TensorType] = None,
-    reduction: int = 1,
-    ignore_index: INT64 = -100,
+@torch_op("aten::cross_entropy")
+def aten_cross_entropy(
+    self: TFloatOrBFloat16,
+    target: Sequence[int],
+    weight: Optional[TFloatOrBFloat16] = None,
+    reduction: str = "mean",
+    ignore_index: int = -100,
     label_smoothing: float = 0.0,
-) -> TensorType:
+) -> 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 op.OptionalHasElement(weight):
+        result, _ = op.SoftmaxCrossEntropyLoss(
+            self, target, reduction=reduction, ignore_index=ignore_index
+        )
+    else:
+        result, _ = op.SoftmaxCrossEntropyLoss(
+            self, target, weight, reduction=reduction, ignore_index=ignore_index
+        )
+    return result
 
 
 @torch_op("aten::elu")

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -184,6 +184,16 @@ 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]]:
+    # change torch.Tensor type to np.array to avoid crash
+    if "weight" in kwargs:
+        if isinstance(kwargs["weight"], torch.Tensor):
+            kwargs["weight"] = np.array(kwargs["weight"])
+    return args, kwargs
+
+
 def _embedding_input_wrangler(
     args: list[Any], kwargs: dict[str, Any]
 ) -> tuple[list[Any], dict[str, Any]]:
@@ -312,6 +322,7 @@ def _where_input_wrangler(
     "nn.functional.adaptive_avg_pool1d": nn_ops.aten_adaptive_avg_pool1d,
     "nn.functional.adaptive_avg_pool2d": nn_ops.aten_adaptive_avg_pool2d,
     "nn.functional.adaptive_avg_pool3d": nn_ops.aten_adaptive_avg_pool3d,
+    "nn.functional.cross_entropy": (nn_ops.aten_cross_entropy, _cross_entropy_input_wrangler),
     "nn.functional.celu": nn_ops.aten_celu,
     "nn.functional.elu": nn_ops.aten_elu,
     "nn.functional.embedding": (core_ops.aten_embedding, _embedding_input_wrangler),
@@ -481,6 +492,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: sample.args[0].dtype != torch.int,
+        reason="ort can only accept int value for the [target] parameter",
+    ),
     skip(
         "nn.functional.upsample_nearest2d",
         # Shape should be [N, C, H, W]