Repository: pytorch/functorch (language: python)

Issue:
Implement batching for nll_loss through decomposition
It assumes that the input has size (N, C) and target has size (C,). There is an edge case where the input may be of size (C,) and the target of size ().


Discussion / hints:
@zou3519 do you mean to support use-cases like:
```python
import torch
from functorch import vmap
from torch.nn.functional import nll_loss

p = torch.rand(4, 5)
t = torch.randint(0, 5, size=(4,))

def f(preds, targets):
    return nll_loss(preds, targets)

out = vmap(f)(p, t)
```
?
Yes, exactly, that gives `RuntimeError: self.dim() == 3 && target.dim() == 2INTERNAL ASSERT FAILED at "/private/home/rzou/functorch/functorch/csrc/BatchRulesLoss.cpp":75, please report a bug to PyTorch`.

nll_loss is a bit tricky because it's difficult to support all of its arguments (weight and ignore_index)
@zou3519 I was thinking that https://github.com/facebookresearch/functorch/blob/8ab5853a303746f36d93fd46e2af4f93d58b6bb9/functorch/csrc/BatchRulesLoss.cpp#L79-L80
and
https://github.com/facebookresearch/functorch/blob/8ab5853a303746f36d93fd46e2af4f93d58b6bb9/functorch/csrc/BatchRulesLoss.cpp#L83
can be optional if `self.dim() == 2` and `target.dim() == 1`. 
That seems reasonable. If `self.dim() == 2 and target.dim() == 1` we probably want to move the batch dim for `self` to the front (otherwise the call to nll_loss won't work)
We talked offline, and we think that it's reasonable for us to just always decompose `nll_loss` to make it easier to write batching rules. `nll_loss` isn't super performance critical, but we should still benchmark it a bit.
To clarify, Chillee is proposing that we:
1) get rid of the special case at https://github.com/facebookresearch/functorch/blob/8ab5853a303746f36d93fd46e2af4f93d58b6bb9/functorch/csrc/BatchRulesLoss.cpp#L139-L147
2) Use the decomposition to cover everything: https://github.com/facebookresearch/functorch/blob/8ab5853a303746f36d93fd46e2af4f93d58b6bb9/functorch/csrc/BatchRulesLoss.cpp#L150-L167 and add in support for `weight` and `ignore_index` to it. We probably need to gather from the weight as well?
@zou3519 thanks for details. I'll work on that, please assign this issue to me.

Fix patch:
diff --git a/functorch/csrc/BatchRulesHelper.h b/functorch/csrc/BatchRulesHelper.h
index 896b1c2b2..373afdf8a 100644
--- a/functorch/csrc/BatchRulesHelper.h
+++ b/functorch/csrc/BatchRulesHelper.h
@@ -392,6 +392,16 @@ inline int64_t get_bdim_size2(
   TORCH_INTERNAL_ASSERT(false);
 }
 
+// [start, start + 1, ..., stop - 1]
+inline VmapDimVector range(int64_t start, int64_t stop) {
+  TORCH_INTERNAL_ASSERT(stop >= start);
+  VmapDimVector dims;
+  dims.reserve(stop - start);
+  for (int64_t i = start; i < stop; i++) {
+    dims.emplace_back(i);
+  }
+  return dims;
+}
 
 }}
 
diff --git a/functorch/csrc/BatchRulesLoss.cpp b/functorch/csrc/BatchRulesLoss.cpp
index b3d342193..77db2a854 100644
--- a/functorch/csrc/BatchRulesLoss.cpp
+++ b/functorch/csrc/BatchRulesLoss.cpp
@@ -67,108 +67,79 @@ mse_loss_backward_batch_rule(
   return std::make_tuple(result, 0);
 };
 
-std::tuple<at::Tensor,optional<int64_t>,at::Tensor,optional<int64_t>>
-nll_loss_forward_self_target_batch_rule(
-    const at::Tensor & self, optional<int64_t> self_bdim,
-    const at::Tensor & target, optional<int64_t> target_bdim,
-    int64_t reduction) {
-  TORCH_INTERNAL_ASSERT(self.dim() == 3 && target.dim() == 2);
-
-  if (reduction == Reduction::None) {
-    int64_t batch_size = self.size(*self_bdim);
-    auto self_ = reshape_dim_into(*self_bdim, 0, self);
-    auto target_ = reshape_dim_into(*target_bdim, 0, target);
-    auto result = at::nll_loss_forward(self_, target_, nullopt, reduction, -100);
-    return std::make_tuple(
-      reshape_dim_outof(0, batch_size, std::get<0>(result)), 0,
-      std::get<1>(result), nullopt
-    );
-  } else if (reduction == Reduction::Sum) {
-    int64_t batch_size = self.size(*self_bdim);
-    auto self_ = reshape_dim_into(*self_bdim, 0, self);
-    auto target_ = reshape_dim_into(*target_bdim, 0, target);
-    auto res = at::nll_loss_forward(self_, target_, nullopt, Reduction::None, -100);
-    auto output = std::get<0>(res);
-    output = reshape_dim_outof(0, batch_size, output);
-    auto total_weight = self_.new_full({}, output.size(-1));
-    return std::make_tuple(
-      output.sum(-1), 0,
-      // NB: total_weight = 0 after Reduction::None
-      total_weight, nullopt
-    );
-  } else if (reduction == Reduction::Mean) {
-    int64_t batch_size = self.size(*self_bdim);
-    auto self_ = reshape_dim_into(*self_bdim, 0, self);
-    auto target_ = reshape_dim_into(*target_bdim, 0, target);
-    auto res = at::nll_loss_forward(self_, target_, nullopt, Reduction::None, -100);
-    auto output = std::get<0>(res);
-    output = reshape_dim_outof(0, batch_size, output);
-    auto total_weight = self_.new_full({}, output.size(-1));
-    return std::make_tuple(
-      output.mean(-1), 0,
-      // NB: total_weight = 0 after Reduction::None
-      total_weight, nullopt
-    );
-  }
-  TORCH_INTERNAL_ASSERT(false);
-}
-
-std::tuple<at::Tensor,at::Tensor> nll_loss_forward_plumbing(
-    const at::Tensor & self,
-    const at::Tensor & target,
-    const c10::optional<at::Tensor> & weight,
+std::tuple<Tensor, Tensor> nll_loss_forward_decomposition(
+    const Tensor & self,
+    const Tensor & target,
+    const c10::optional<Tensor> & weight,
     int64_t reduction, int64_t ignore_index) {
-  auto maybe_layer = maybeCurrentDynamicLayer();
-  TORCH_INTERNAL_ASSERT(maybe_layer.has_value());
-  int64_t cur_level = maybe_layer->layerId();
-
-  Tensor self_value;
-  optional<int64_t> self_bdim;
-  std::tie(self_value, self_bdim) = unwrapTensorAtLevel(self, cur_level);
 
-  Tensor target_value;
-  optional<int64_t> target_bdim;
-  std::tie(target_value, target_bdim) = unwrapTensorAtLevel(target, cur_level);
-
-  optional<Tensor> weight_value;
-  optional<int64_t> weight_bdim;
-  if (weight) {
-    std::tie(weight_value, weight_bdim) = unwrapTensorAtLevel(*weight, cur_level);
+  bool has_ignore_index = ignore_index >= 0;
+  if (has_ignore_index) {
+    // fallback
+    if (target.dim() > 1) {
+      static auto op = c10::Dispatcher::singleton()
+        .findSchemaOrThrow("aten::nll_loss_nd", "");
+      return slow_fallback<Tensor, Tensor>(op, {self, target, weight, reduction, ignore_index});
+    } else {
+      static auto op = c10::Dispatcher::singleton()
+        .findSchemaOrThrow("aten::nll_loss_forward", "");
+      return slow_fallback<Tensor, Tensor>(op, {self, target, weight, reduction, ignore_index});
+    }
   }
+  // self can be [N, C, ...] or [C]
+  // target can be [N, ...] or []
 
-  if (self_bdim && target_bdim && (!weight || !weight->defined()) && ignore_index < 0) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    auto results = nll_loss_forward_self_target_batch_rule(
-        self_value, self_bdim, target_value, target_bdim, reduction);
-    return std::make_tuple(
-      makeBatched(std::get<0>(results), std::get<1>(results), cur_level),
-      makeBatched(std::get<2>(results), std::get<3>(results), cur_level)
-    );
+  int64_t channel_dim = 1;
+  if (self.dim() < 2) {
+    channel_dim = 0;
   }
-
-  if ((!weight || !weight->defined()) && ignore_index < 0) {
-    // Decomposition: gather to get unreduced loss. 1 is for the C dim, that's always 1.
-    // gather can handle arbitrary strides so it's a good candidate for a decomposition.
-    auto target_ = target.unsqueeze(1);
-    auto result = at::gather(self, 1, target_).squeeze(1);
-    auto total_weight = at::full(
-        {}, result.numel(), self.scalar_type(),
-        self.layout(), self.device(), nullopt);
-
-    // Apply the reduction
-    switch (reduction) {
-      case Reduction::None:
-        return std::make_tuple(-result, total_weight);
-      case Reduction::Sum:
-        return std::make_tuple(-result.sum(), total_weight);
-      case Reduction::Mean:
-        return std::make_tuple(-result.mean(), total_weight);
+  auto self_ = self;
+  Tensor weight_;
+
+  if (weight && weight->defined()) {
+    // Here is a specific case with reduction mean and non-batched tensors
+    // https://github.com/pytorch/pytorch/issues/61309
+    // In this case weight is cancelled: w * x[t] / w -> x[t]
+    if (!(reduction == Reduction::Mean && self_.dim() < 2)) {
+      // reshape weights to [1, C, 1, ..., 1]
+      auto shape = weight->sizes();
+      VmapDimVector new_shape(self_.dim(), 1);
+      new_shape[channel_dim] = shape[0];
+      weight_ = weight->reshape(new_shape);
+      self_ = self_ * weight_;
     }
   }
+  auto target_ = target.unsqueeze(channel_dim);
+  // target can be [N, 1, ...] or [1]
+
+  auto result = -at::gather(self_, channel_dim, target_).squeeze(channel_dim);
+  auto total_weight = at::full(
+      {}, result.numel(), self_.scalar_type(),
+      self_.layout(), self_.device(), nullopt);
+
+  // Apply the reduction
+  if (result.dim() > 0) {
+    if (reduction == Reduction::Sum) {
+      result = result.sum();
+    } else if (reduction == Reduction::Mean) {
+      if (!weight || !weight->defined()) {
+        result = result.mean();
+      } else {
+        TORCH_INTERNAL_ASSERT(weight_.defined());
+        weight_ = weight_.expand(self_.sizes());
+        auto wsum = at::gather(weight_, channel_dim, target_).squeeze(channel_dim);
+        wsum = wsum.sum();
+        result = result.sum() / wsum;
+        total_weight = wsum;
+      }
+    }
+  } else if (reduction == Reduction::Mean && weight && weight->defined()) {
+    // here weight is [C] and target is [1]
+    auto wsum = at::gather(*weight, channel_dim, target_).squeeze(channel_dim);
+    total_weight = wsum.sum();
+  }
 
-  static auto op = c10::Dispatcher::singleton()
-    .findSchemaOrThrow("aten::nll_loss_forward", "");
-  return slow_fallback<Tensor,Tensor>(op, {self, target, weight, reduction, ignore_index});
+  return std::make_tuple(result, total_weight);
 }
 
 std::tuple<at::Tensor,optional<int64_t>>
@@ -272,7 +243,7 @@ at::Tensor nll_loss_backward_plumbing(
 
 
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
-  m.impl("nll_loss_forward", nll_loss_forward_plumbing);
+  m.impl("nll_loss_forward", nll_loss_forward_decomposition);
   m.impl("nll_loss_backward", nll_loss_backward_plumbing);
   VMAP_SUPPORT("mse_loss", mse_loss_batch_rule);
   VMAP_SUPPORT("mse_loss_backward", mse_loss_backward_batch_rule);
diff --git a/functorch/csrc/BatchRulesReduceOps.cpp b/functorch/csrc/BatchRulesReduceOps.cpp
index 5df9acaba..b3067e012 100644
--- a/functorch/csrc/BatchRulesReduceOps.cpp
+++ b/functorch/csrc/BatchRulesReduceOps.cpp
@@ -11,18 +11,6 @@
 
 namespace at { namespace functorch {
 
-// [start, start + 1, ..., stop - 1]
-static VmapDimVector range(int64_t start, int64_t stop) {
-  TORCH_INTERNAL_ASSERT(stop >= start);
-  VmapDimVector dims;
-  dims.reserve(stop - start);
-  for (int64_t i = start; i < stop; i++) {
-    dims.emplace_back(i);
-  }
-  return dims;
-}
-
-
 bool is_allowed_dim_on_scalar_tensor(int64_t dim) {
   return dim == 0 || dim == -1;
 }


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
index 9a3f384f2..d30107055 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -360,7 +360,6 @@ def vjp_of_vjp(*args_and_cotangents):
         xfail('nanmean'),
         xfail('block_diag'),
         xfail('nn.functional.dropout'),
-        xfail('nn.functional.nll_loss'),
     }))
     def test_vmapvjp(self, device, dtype, op):
         # These are too annoying to put into the list above
@@ -539,7 +538,6 @@ def test():
         xfail('vstack'),
         xfail('block_diag'),
         xfail('nn.functional.batch_norm'),
-        xfail('nn.functional.nll_loss'),
         xfail('cdist'),
         xfail('lu_solve'),
         xfail('lu_unpack'),
diff --git a/test/test_vmap.py b/test/test_vmap.py
index ab4b51252..4d732b098 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -3015,7 +3015,6 @@ class TestVmapOperatorsOpInfo(TestCase):
         xfail('fft.rfft'),
         xfail('fft.rfftn'),
         xfail('nn.functional.batch_norm'),
-        xfail('nn.functional.nll_loss'),
         xfail('lu_unpack'),
         xfail('nn.functional.pad', 'constant'),
     })