feat: align nll_loss with PyTorch Python API

src/base/nll_loss.h

@@ -2,13 +2,39 @@
 #define INFINI_OPS_BASE_NLL_LOSS_H_
 
 #include <optional>
+#include <string>
 
 #include "operator.h"
 
 namespace infini::ops {
 
 class NllLoss : public Operator<NllLoss> {
  public:
+  NllLoss(const Tensor input, const Tensor target,
+          const std::optional<Tensor> weight,
+          const std::optional<bool> size_average, const int64_t ignore_index,
+          const std::optional<bool> reduce, const std::string reduction,
+          Tensor out)
+      : input_shape_{input.shape()},
+        input_strides_{input.strides()},
+        input_type_{input.dtype()},
+        target_shape_{target.shape()},
+        target_strides_{target.strides()},
+        target_type_{target.dtype()},
+        out_shape_{out.shape()},
+        out_strides_{out.strides()},
+        out_type_{out.dtype()},
+        has_weight_{weight.has_value()},
+        weight_shape_{weight ? weight->shape() : Tensor::Shape{}},
+        weight_strides_{weight ? weight->strides() : Tensor::Strides{}},
+        weight_type_{weight ? weight->dtype() : DataType::kFloat32},
+        reduction_{
+            ReductionFromPythonArguments(size_average, reduce, reduction)},
+        ignore_index_{ignore_index},
+        device_index_{out.device().index()} {}
+
+  [[deprecated(
+      "Use the overload with `ignore_index` before string `reduction`.")]]
   NllLoss(const Tensor input, const Tensor target,
           const std::optional<Tensor> weight, const int64_t reduction,
           const int64_t ignore_index, Tensor out)
@@ -29,6 +55,19 @@ class NllLoss : public Operator<NllLoss> {
         ignore_index_{ignore_index},
         device_index_{out.device().index()} {}
 
+  void operator()(const Tensor input, const Tensor target,
+                  const std::optional<Tensor> weight,
+                  const std::optional<bool> size_average,
+                  const int64_t ignore_index, const std::optional<bool> reduce,
+                  const std::string reduction, Tensor out) const {
+    return operator()(
+        input, target, weight,
+        ReductionFromPythonArguments(size_average, reduce, reduction),
+        ignore_index, out);
+  }
+
+  [[deprecated(
+      "Use the overload with `ignore_index` before string `reduction`.")]]
   virtual void operator()(const Tensor input, const Tensor target,
                           const std::optional<Tensor> weight,
                           const int64_t reduction, const int64_t ignore_index,
@@ -66,6 +105,40 @@ class NllLoss : public Operator<NllLoss> {
   int64_t ignore_index_{};
 
   int device_index_{0};
+
+ private:
+  static int64_t ReductionFromPythonArguments(
+      const std::optional<bool> size_average, const std::optional<bool> reduce,
+      const std::string& reduction) {
+    if (!size_average.has_value() && !reduce.has_value()) {
+      return ReductionFromString(reduction);
+    }
+
+    if (!reduce.value_or(true)) {
+      return 0;
+    }
+
+    if (!size_average.value_or(true)) {
+      return 2;
+    }
+
+    return 1;
+  }
+
+  static int64_t ReductionFromString(const std::string& reduction) {
+    if (reduction == "none") {
+      return 0;
+    }
+
+    if (reduction == "mean") {
+      return 1;
+    }
+
+    assert(reduction == "sum" &&
+           "`NllLoss` reduction must be `none`, `mean`, or `sum`");
+
+    return 2;
+  }
 };
 
 }  // namespace infini::ops

tests/test_nll_loss.py

@@ -0,0 +1,165 @@
+import infini.ops
+import pytest
+
+import torch
+from tests.utils import (
+    Payload,
+    empty_strided,
+    get_stream,
+    rand_strided,
+    randint_strided,
+    randn_strided,
+)
+
+
+@pytest.mark.auto_act_and_assert
+@pytest.mark.parametrize(
+    "expected_reduction, has_weight, ignore_index, api_style, size_average, reduce, reduction",
+    (
+        ("mean", False, -100, "keyword", None, None, "mean"),
+        ("none", False, -100, "full", None, None, "none"),
+        ("mean", True, -100, "full", None, None, "mean"),
+        ("sum", True, 2, "legacy", None, None, "sum"),
+        ("sum", False, -100, "full", None, None, "sum"),
+        ("none", False, -100, "full", None, False, "sum"),
+        ("sum", True, -100, "full", False, None, "none"),
+        ("mean", True, -100, "full", None, True, "sum"),
+    ),
+)
+@pytest.mark.parametrize(
+    "dtype, rtol, atol",
+    (
+        (torch.float32, 1e-5, 1e-5),
+        (torch.float16, 1e-2, 1e-2),
+        (torch.bfloat16, 1e-2, 1e-2),
+    ),
+)
+def test_nll_loss(
+    expected_reduction,
+    has_weight,
+    ignore_index,
+    api_style,
+    size_average,
+    reduce,
+    reduction,
+    implementation_index,
+    dtype,
+    device,
+    rtol,
+    atol,
+):
+    batch_size = 7
+    num_classes = 5
+    input = torch.nn.functional.log_softmax(
+        randn_strided((batch_size, num_classes), None, dtype=dtype, device=device),
+        dim=1,
+    )
+    target = randint_strided(
+        0,
+        num_classes,
+        (batch_size,),
+        None,
+        dtype=torch.int64,
+        device=device,
+    )
+
+    if ignore_index >= 0:
+        target[0] = ignore_index
+
+    weight = None
+
+    if has_weight:
+        weight = rand_strided((num_classes,), None, dtype=dtype, device=device).add_(
+            0.5
+        )
+
+    out_shape = target.shape if expected_reduction == "none" else ()
+    out = empty_strided(out_shape, None, dtype=dtype, device=device)
+
+    return Payload(
+        lambda *args: _nll_loss(
+            *args,
+            reduction=reduction,
+            ignore_index=ignore_index,
+            api_style=api_style,
+            size_average=size_average,
+            reduce=reduce,
+            implementation_index=implementation_index,
+        ),
+        lambda *args: _torch_nll_loss(
+            *args, reduction=expected_reduction, ignore_index=ignore_index
+        ),
+        (input, target, weight, out),
+        {},
+        rtol=rtol,
+        atol=atol,
+    )
+
+
+def _nll_loss(
+    input,
+    target,
+    weight,
+    out,
+    *,
+    reduction,
+    ignore_index,
+    api_style,
+    size_average,
+    reduce,
+    implementation_index,
+):
+    kwargs = {
+        "implementation_index": implementation_index,
+        "stream": get_stream(input.device),
+    }
+
+    if api_style == "keyword":
+        infini.ops.nll_loss(
+            input=input,
+            target=target,
+            weight=weight,
+            size_average=size_average,
+            ignore_index=ignore_index,
+            reduce=reduce,
+            reduction=reduction,
+            out=out,
+            **kwargs,
+        )
+    elif api_style == "legacy":
+        infini.ops.nll_loss(
+            input,
+            target,
+            weight,
+            {"none": 0, "mean": 1, "sum": 2}[reduction],
+            ignore_index,
+            out,
+            **kwargs,
+        )
+    else:
+        infini.ops.nll_loss(
+            input,
+            target,
+            weight,
+            size_average,
+            ignore_index,
+            reduce,
+            reduction,
+            out,
+            **kwargs,
+        )
+
+    return out
+
+
+def _torch_nll_loss(input, target, weight, out, *, reduction, ignore_index):
+    result = torch.nn.functional.nll_loss(
+        input,
+        target,
+        weight=weight,
+        ignore_index=ignore_index,
+        reduction=reduction,
+    )
+    out.copy_(result)
+
+    return out