Repository: pytorch/functorch (language: python)

Issue:
Can't implement batching rule for F.one_hot as written
We can't currently implement a batching rule for `F.one_hot`, even if we had batching rules for the constituent ops (`_local_scalar_dense`, `min`, `max`, `unsqueeze`, `scatter`).

In the C++ code, the code looks like this here (https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Onehot.cpp#L38)
```
    Tensor ret = at::zeros(shape, self.options());
    ret.scatter_(-1, self.unsqueeze(-1), 1);
    return ret;
```
cc: @bdhirsh



Discussion / hints:
Any ideas how to solve this? I've been using vmap and found this severely impacted my performance.
@pedro-abreu We can prioritize fixing it - there's a couple different ways we can resolve it.

One of them is by decomposing it into a `zeros` + non-inplace scatter. We can do that.

Fix patch:
diff --git a/functorch/csrc/BatchRulesModules.cpp b/functorch/csrc/BatchRulesModules.cpp
index 662857668..b9b978da1 100644
--- a/functorch/csrc/BatchRulesModules.cpp
+++ b/functorch/csrc/BatchRulesModules.cpp
@@ -200,6 +200,42 @@ std::tuple<Tensor,Tensor> cudnn_convolution_backward_plumbing(const Tensor & sel
   return slow_fallback<Tensor,Tensor>(op, { self, grad_output, weight, padding, stride, dilation, groups, benchmark, deterministic, allow_tf32, output_mask });
 }
 
+// TODO: replace with targetable functionalization
+Tensor one_hot_decomposition_hack(const Tensor &self, int64_t num_classes) {
+    TORCH_CHECK(self.dtype() == kLong, "one_hot is only applicable to index tensor.");
+    auto shape = self.sizes().vec();
+
+    // empty tensor could be converted to one hot representation,
+    // but shape inference is not possible.
+    if (self.numel() == 0) {
+        if (num_classes <= 0) {
+            AT_ERROR("Can not infer total number of classes from empty tensor.");
+        } else {
+            shape.push_back(num_classes);
+            return at::empty(shape, self.options());
+        }
+    }
+
+    TORCH_CHECK(num_classes > 0, "When vmap-ing torch.nn.functional.one_hot, please "
+        "provide an explicit positive num_classes argument.");
+
+    // Disabling all of the following checks. This is OK because scatter has checks too.
+    // Maybe one_hot should be a primitive wrt autograd so we don't have to deal with this.
+    // // non-empty tensor
+    // if (self.device().type() != at::kCUDA) {
+    //   //for cuda, rely on device assert thrown by scatter
+    //   TORCH_CHECK(self.min().item().toLong() >= 0, "Class values must be non-negative.");
+    // }
+    // if (self.device().type() != at::kCUDA) {
+    //   //rely on device asserts from scatter to avoid sync here
+    //   TORCH_CHECK(num_classes > self.max().item().toLong(), "Class values must be smaller than num_classes.");
+    // }
+
+    shape.push_back(num_classes);
+    Tensor ret = at::zeros(shape, self.options());
+    return ret.scatter(-1, self.unsqueeze(-1), 1);
+}
+
 template <typename A, A a, typename C>
 struct UpsampleBackwardBatchRuleHelper;
 
@@ -264,6 +300,6 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   UPSAMPLE_BACKWARD(upsample_nearest2d_backward, vec);
   UPSAMPLE_BACKWARD(upsample_nearest3d_backward, vec);
   UPSAMPLE_BACKWARD(upsample_trilinear3d_backward, vec);
-
+  m.impl("one_hot", one_hot_decomposition_hack);
 }
 }}
diff --git a/functorch/csrc/BatchRulesScatterOps.cpp b/functorch/csrc/BatchRulesScatterOps.cpp
index b6e76b4ab..d9367f62b 100644
--- a/functorch/csrc/BatchRulesScatterOps.cpp
+++ b/functorch/csrc/BatchRulesScatterOps.cpp
@@ -97,6 +97,72 @@ int64_t bdim_size(
   TORCH_INTERNAL_ASSERT(false);
 }
 
+std::tuple<Tensor,optional<int64_t>> scatter_value_batch_rule(
+    const Tensor& self, optional<int64_t> self_bdim,
+    int64_t dim,
+    const Tensor& index, optional<int64_t> index_bdim,
+    const Scalar& value) {
+  auto self_logical_rank = rankWithoutBatchDim(self, self_bdim);
+  auto index_logical_rank = rankWithoutBatchDim(index, index_bdim);
+  auto batch_size = bdim_size(self, self_bdim, index, index_bdim);
+
+  auto self_ = moveBatchDimToFront(self, self_bdim);
+  auto index_ = moveBatchDimToFront(index, index_bdim);
+
+  if (self_logical_rank == 0) {
+    self_ = self_.unsqueeze(-1);
+  }
+  if (index_logical_rank == 0) {
+    index_ = index_.unsqueeze(-1);
+  }
+  self_ = ensure_has_bdim(self_, self_bdim.has_value(), batch_size);
+  index_ = ensure_has_bdim(index_, index_bdim.has_value(), batch_size);
+  auto physical_dim = getPhysicalDim(self_, /*has_batch_dim*/true, dim);
+
+  auto result = at::scatter(self_, physical_dim, index_, value);
+  // result should have same shape as self
+  if (self_logical_rank == 0) {
+    result = result.squeeze(-1);
+  }
+  return std::make_tuple(result, 0);
+}
+
+std::tuple<Tensor,optional<int64_t>> scatter_src_batch_rule(
+    const Tensor& self, optional<int64_t> self_bdim,
+    int64_t dim,
+    const Tensor& index, optional<int64_t> index_bdim,
+    const Tensor& src, optional<int64_t> src_bdim) {
+  auto self_logical_rank = rankWithoutBatchDim(self, self_bdim);
+  auto index_logical_rank = rankWithoutBatchDim(index, index_bdim);
+  auto src_logical_rank = rankWithoutBatchDim(src, src_bdim);
+  auto batch_size = bdim_size(self, self_bdim, index, index_bdim, src, src_bdim);
+
+  auto self_ = moveBatchDimToFront(self, self_bdim);
+  auto index_ = moveBatchDimToFront(index, index_bdim);
+  auto src_ = moveBatchDimToFront(src, src_bdim);
+
+  if (self_logical_rank == 0) {
+    self_ = self_.unsqueeze(-1);
+  }
+  if (index_logical_rank == 0) {
+    index_ = index_.unsqueeze(-1);
+  }
+  if (src_logical_rank == 0) {
+    src_ = src_.unsqueeze(-1);
+  }
+  self_ = ensure_has_bdim(self_, self_bdim.has_value(), batch_size);
+  index_ = ensure_has_bdim(index_, index_bdim.has_value(), batch_size);
+  src_ = ensure_has_bdim(src_, src_bdim.has_value(), batch_size);
+  auto physical_dim = getPhysicalDim(self_, /*has_batch_dim*/true, dim);
+
+  auto result = at::scatter(self_, physical_dim, index_, src_);
+  // result should have same shape as self
+  if (self_logical_rank == 0) {
+    result = result.squeeze(-1);
+  }
+  return std::make_tuple(result, 0);
+}
+
 std::tuple<Tensor,optional<int64_t>> gather_batch_rule(
     const Tensor& self, optional<int64_t> self_bdim,
     int64_t dim,
@@ -168,6 +234,8 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
     // m.impl("index.Tensor", index_plumbing);
   VMAP_SUPPORT("gather", gather_batch_rule);
   VMAP_SUPPORT("gather_backward", gather_backward_batch_rule);
+  VMAP_SUPPORT("scatter.value", scatter_value_batch_rule);
+  VMAP_SUPPORT("scatter.src", scatter_src_batch_rule);
 }
 
 }}


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 24627e7f9..85b367be1 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2542,6 +2542,14 @@ def test_conv2d(self):
             for loop_out, batched_out in get_fallback_and_vmap_exhaustive(conv_fn, arg_values, kwarg_values):
                 self.assertEqual(loop_out, batched_out)
 
+    def test_one_hot(self):
+        sample_inputs = [
+            (torch.randint(0, 3, []), 3),
+            (torch.randint(0, 3, [2, 3, 4]), 4),
+        ]
+        for args in sample_inputs:
+            for loop_out, batched_out in get_fallback_and_vmap_exhaustive(F.one_hot, args, {}):
+                self.assertEqual(loop_out, batched_out)
 
     def test_mode_key(self):
         def vmap_f(x):