Ensure tensors devices match on torch.index_put batch rule impl (pytorch#130479)

Pull Request resolved: pytorch#130479
Approved by: https://github.com/zou3519

Author: guilhermeleobas

aten/src/ATen/functorch/BatchRulesScatterOps.cpp

@@ -498,11 +498,18 @@ Tensor& index_put__plumbing(Tensor & self, const List<optional<Tensor>> & indice
   auto maybe_layer = maybeCurrentDynamicLayer();
   vmap_check_escaped(maybe_layer, "index_put__plumbing");
   int64_t cur_level = maybe_layer->layerId();
-  if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
-    return self.index_put_(indices, values, accumulate);
+
+  // on device mismatch, we can move 0d tensors to self device
+  auto values_ = values;
+  if (values.device() != self.device() && values.numel() == 1 && values.dim() == 0) {
+    values_ = values.to(self.device());
+  }
+
+  if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values_, cur_level)) {
+    return self.index_put_(indices, values_, accumulate);
   }
   auto [self_value, self_bdim, indices_value, indices_bdims, values_value, values_bdim] =
-      unpackSelfAndIndicesAndValuesAtCurrentLevel(self, indices, values, cur_level);
+      unpackSelfAndIndicesAndValuesAtCurrentLevel(self, indices, values_, cur_level);
   index_put__batch_rule(self_value, self_bdim, indices_value, indices_bdims, values_value, values_bdim, accumulate);
   return self;
 }
@@ -645,11 +652,18 @@ Tensor index_put_plumbing(const Tensor & self, const List<optional<Tensor>> & in
   auto maybe_layer = maybeCurrentDynamicLayer();
   vmap_check_escaped(maybe_layer, "index_put_plumbing");
   int64_t cur_level = maybe_layer->layerId();
-  if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values, cur_level)) {
-    return self.index_put(indices, values, accumulate);
+
+  // on device mismatch, we can move 0d tensors to self device
+  auto values_ = values;
+  if (values.device() != self.device() && values.numel() == 1 && values.dim() == 0) {
+    values_ = values.to(self.device());
+  }
+
+  if (!isBatchedAtLevel(self, cur_level) && !isBatchedAtLevel(indices, cur_level) && !isBatchedAtLevel(values_, cur_level)) {
+    return self.index_put(indices, values_, accumulate);
   }
   auto [self_value, self_bdim, indices_value, indices_bdims, values_value, values_bdim] =
-      unpackSelfAndIndicesAndValuesAtCurrentLevel(self, indices, values, cur_level);
+      unpackSelfAndIndicesAndValuesAtCurrentLevel(self, indices, values_, cur_level);
   auto results = index_put_batch_rule(self_value, self_bdim, indices_value, indices_bdims, values_value, values_bdim, accumulate);
   return makeBatched(std::get<0>(results), std::get<1>(results), cur_level);
 }

test/functorch/test_vmap.py

@@ -49,6 +49,7 @@
 from torch.testing._internal.common_cuda import with_tf32_off
 from torch.testing._internal.common_device_type import (
     instantiate_device_type_tests,
+    onlyCUDA,
     OpDTypes,
     ops,
     tol,
@@ -4793,6 +4794,21 @@ def f(x, gy):
 
         self.vmap_outplace_test(f, (x, gy), {}, in_dims=(None, 0))
 
+    @onlyCUDA
+    @parametrize("inplace", [True, False])
+    def test_0d_tensor_index_put(self, device, inplace):
+        def f(t, idx, v):
+            fn = torch.index_put_ if inplace else torch.index_put
+            return fn(t, idx, v)
+
+        N = 2
+        t = torch.zeros((N, 5), device="cuda")
+        idx = torch.tensor([1, 3])
+        v = torch.tensor(1, dtype=t.dtype, device="cpu")
+
+        expected = torch.tensor([[0, 1, 0, 1, 0], [0, 1, 0, 1, 0]], dtype=t.dtype)
+        self.assertEqual(expected, vmap(f, in_dims=(0, None, None))(t, (idx,), v))
+
     @parametrize("training", [True, False])
     @parametrize("track_running_stats", [True, False])
     @parametrize("affine", [True, False])