feat(atenlib): add ops(nll_loss)

onnxscript/function_libs/torch_aten/ops/nn.py

@@ -769,16 +769,77 @@ def aten_multilabel_margin_loss_forward(
     raise NotImplementedError()
 
 
+@torch_op("aten::nll_loss")
 def aten_nll_loss(
-    self: TensorType,
-    target: TensorType,
-    weight: Optional[TensorType] = None,
+    self: TFloat,
+    target: INT64,
     reduction: int = 1,
-    ignore_index: INT64 = -100,
-) -> TensorType:
+    ignore_index: int = -100,
+) -> TFloat:
     """nll_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100) -> Tensor"""
 
-    raise NotImplementedError()
+    rank_self = op.Size(op.Shape(self))
+    if rank_self == 1:  # self rank should be at least 2
+        self = op.Unsqueeze(self, op.Constant(value_ints=[0]))
+
+    rank_target = op.Size(op.Shape(target))
+    if rank_target == 0:  # target rank should be at least 1
+        target = op.Unsqueeze(target, op.Constant(value_ints=[0]))
+
+    if reduction == 0:
+        result = op.NegativeLogLikelihoodLoss(
+            self, target, ignore_index=ignore_index, reduction="none"
+        )
+    elif reduction == 1:
+        result = op.NegativeLogLikelihoodLoss(
+            self, target, ignore_index=ignore_index, reduction="mean"
+        )
+    else:  # assert reduction == 2
+        result = op.NegativeLogLikelihoodLoss(
+            self, target, ignore_index=ignore_index, reduction="sum"
+        )
+
+    if rank_self == 1:
+        result = op.Squeeze(result)
+
+    return result
+
+
+@torch_op("aten::nll_loss", overload=True)
+def aten_nll_loss_weight(
+    self: TFloat,
+    target: INT64,
+    weight: TFloat,
+    reduction: int = 1,
+    ignore_index: int = -100,
+) -> TFloat:
+    """nll_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100) -> Tensor"""
+
+    rank_self = op.Size(op.Shape(self))
+    if rank_self == 1:  # self rank should be at least 2
+        self = op.Unsqueeze(self, op.Constant(value_ints=[0]))
+
+    rank_target = op.Size(op.Shape(target))
+    if rank_target == 0:  # target rank should be at least 1
+        target = op.Unsqueeze(target, op.Constant(value_ints=[0]))
+
+    if reduction == 0:
+        result = op.NegativeLogLikelihoodLoss(
+            self, target, weight, ignore_index=ignore_index, reduction="none"
+        )
+    elif reduction == 1:
+        result = op.NegativeLogLikelihoodLoss(
+            self, target, weight, ignore_index=ignore_index, reduction="mean"
+        )
+    else:
+        result = op.NegativeLogLikelihoodLoss(
+            self, target, weight, ignore_index=ignore_index, reduction="sum"
+        )
+
+    if rank_self == 1:
+        result = op.Squeeze(result)
+
+    return result
 
 
 def aten_nll_loss2d(

onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py

@@ -294,6 +294,17 @@ def _full_input_wrangler(
     return args, kwargs
 
 
+def _nll_loss_input_wrangler(
+    args: list[Any], kwargs: dict[str, Any]
+) -> tuple[list[Any], dict[str, Any]]:
+    if "reduction" in kwargs:
+        # aten_nll_loss can only accept integer argument instead of string
+        reduction_vals = ["none", "mean", "sum"]
+        value = kwargs["reduction"]
+        kwargs["reduction"] = reduction_vals.index(value)
+    return args, kwargs
+
+
 def _upsample_input_wrangler(
     args: list[Any], kwargs: dict[str, Any]
 ) -> tuple[list[Any], dict[str, Any]]:
@@ -429,6 +440,8 @@ def _where_input_wrangler(
     "nn.functional.embedding": (core_ops.aten_embedding, _embedding_input_wrangler),
     "nn.functional.leaky_relu": nn_ops.aten_leaky_relu,
     "nn.functional.logsigmoid": nn_ops.aten_log_sigmoid,
+    "nn.functional.nll_loss_weight": (nn_ops.aten_nll_loss_weight, _nll_loss_input_wrangler),
+    "nn.functional.nll_loss": (nn_ops.aten_nll_loss, _nll_loss_input_wrangler),
     "nn.functional.relu": nn_ops.aten_relu,
     "nn.functional.relu6": nn_ops.aten_relu6,
     "nn.functional.selu": core_ops.aten_selu,
@@ -752,6 +765,16 @@ def _where_input_wrangler(
         matcher=lambda sample: len(sample.kwargs) == 0 or sample.kwargs.get("p", 0.0) > 0.0,
         reason="dropout is random so the result not match",
     ),
+    skip(
+        "nn.functional.nll_loss",
+        matcher=lambda sample: "weight" in sample.kwargs,
+        reason="this Aten overload doesn't accept weight as kwargs",
+    ),
+    skip(
+        "nn.functional.nll_loss_weight",
+        matcher=lambda sample: "weight" not in sample.kwargs,
+        reason="this Aten overload need weight as kwargs",
+    ),
     skip(
         "nn.functional.upsample_nearest2d",
         # Shape should be [N, C, H, W]
@@ -796,6 +819,8 @@ def _where_input_wrangler(
     ),
 )
 
+duplicate_opinfo(OPS_DB, "nn.functional.nll_loss", ("nn.functional.nll_loss_weight",))
+
 duplicate_opinfo(
     OPS_DB,
     "min",
@@ -878,6 +903,8 @@ def _convert_kwargs_for_onnx(kwargs: dict[str, Any]) -> dict[str, Any]:
             continue
         if key == "dtype":
             value = TORCH_TYPE_TO_ONNX[value]
+        if isinstance(value, torch.Tensor):
+            value = np.array(value)
         new_kwargs[key] = value
     return new_kwargs