Fix functional::smooth_l1_loss signatures to not override beta

kurtamohler · kurtamohler · commit 38edd6601ff8 · 2023-10-09T22:14:53.000Z
diff --git a/test/cpp/api/functional.cpp b/test/cpp/api/functional.cpp
@@ -303,6 +303,23 @@ TEST_F(FunctionalTest, SmoothL1LossBeta) {
   ASSERT_TRUE(input.sizes() == input.grad().sizes());
 }
 
+TEST_F(FunctionalTest, SmoothL1LossBetaOptions) {
+  auto input = torch::tensor(
+      {0.1, 1.5, 10.0}, torch::dtype(torch::kFloat).requires_grad(true));
+  auto target = torch::tensor({0., 1., 5.}, torch::kFloat);
+  auto output =
+      // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers)
+      F::smooth_l1_loss(
+          input,
+          target,
+          F::SmoothL1LossFuncOptions().reduction(torch::kMean).beta(0.5));
+  auto expected = torch::tensor(1.67, torch::kFloat);
+  auto s = output.sum();
+  s.backward();
+  ASSERT_TRUE(output.allclose(expected));
+  ASSERT_TRUE(input.sizes() == input.grad().sizes());
+}
+
 TEST_F(FunctionalTest, SmoothL1LossNoReduction) {
   auto input = torch::tensor(
       {0.1, 1.2, 4.7}, torch::dtype(torch::kFloat).requires_grad(true));
diff --git a/torch/csrc/api/include/torch/nn/functional/loss.h b/torch/csrc/api/include/torch/nn/functional/loss.h
@@ -346,7 +346,7 @@ inline Tensor smooth_l1_loss(
     const Tensor& input,
     const Tensor& target,
     SmoothL1LossFuncOptions::reduction_t reduction,
-    double beta = 1.) {
+    c10::optional<double> beta_opt = c10::nullopt) {
   if (target.sizes() != input.sizes()) {
     TORCH_WARN(
         "Using a target size (",
@@ -357,6 +357,7 @@ inline Tensor smooth_l1_loss(
         "This will likely lead to incorrect results due to broadcasting. ",
         "Please ensure they have the same size.");
   }
+  double beta = beta_opt.value_or(1.0);
 
   std::vector<Tensor> expanded_tensors =
       torch::broadcast_tensors({input, target});
@@ -384,8 +385,29 @@ inline Tensor smooth_l1_loss(
 inline Tensor smooth_l1_loss(
     const Tensor& input,
     const Tensor& target,
-    const SmoothL1LossFuncOptions& options = {},
-    double beta = 1.) {
+    const SmoothL1LossFuncOptions& options = {}) {
+  return detail::smooth_l1_loss(
+      input, target, options.reduction(), options.beta());
+}
+
+/// See
+/// https://pytorch.org/docs/master/nn.functional.html#torch.nn.functional.smooth_l1_loss
+/// about the exact behavior of this functional.
+///
+/// Example:
+/// ```
+/// namespace F = torch::nn::functional;
+/// F::smooth_l1_loss(input, target, /*options=*/torch::kNone, /*beta=*/0.5);
+/// ```
+inline Tensor smooth_l1_loss(
+    const Tensor& input,
+    const Tensor& target,
+    const SmoothL1LossFuncOptions& options,
+    double beta) {
+  TORCH_CHECK(
+      options.beta() == c10::nullopt,
+      "expected beta not to be provided in 'options', but got ",
+      options.beta().value());
   return detail::smooth_l1_loss(input, target, options.reduction(), beta);
 }
 
diff --git a/torch/csrc/api/include/torch/nn/options/loss.h b/torch/csrc/api/include/torch/nn/options/loss.h
@@ -546,8 +546,9 @@ struct TORCH_API SmoothL1LossOptions {
   /// be summed. Default: 'mean'
   TORCH_ARG(reduction_t, reduction) = torch::kMean;
   /// Specifies the threshold at which to change between L1 and L2 loss.
-  /// Default: 1.0
-  TORCH_ARG(double, beta) = 1.0;
+  /// If beta is not specified, a value of 1.0 will be used.
+  /// Default: nullopt
+  TORCH_ARG(c10::optional<double>, beta) = c10::nullopt;
 };
 
 namespace functional {
