don't call .item in onehot for XLA (#127335)

JackCaoG · pytorchmergebot · commit 38a33c3202a7 · 2024-05-29T20:37:26.000Z
We found that `nn.function.one_hot` will cause a graph break due to the item call in the native implementation. Pull Request resolved: #127335 Approved by: https://github.com/ezyang
diff --git a/aten/src/ATen/native/Onehot.cpp b/aten/src/ATen/native/Onehot.cpp
@@ -43,15 +43,15 @@ Tensor one_hot(const Tensor &self, int64_t num_classes) {
 
     // non-empty tensor
     if (self.device().type() != at::kCUDA && self.device().type() != at::kMPS &&
-        self.device().type() != at::kPrivateUse1) {
+        self.device().type() != at::kPrivateUse1 && self.device().type() != at::kXLA) {
       // for cuda, rely on device assert thrown by scatter
       TORCH_CHECK(self.min().item().toLong() >= 0, "Class values must be non-negative.");
     }
     if (num_classes == -1) {
         num_classes = self.max().item().toLong() + 1;
     } else {
         if (self.device().type() != at::kCUDA && self.device().type() != at::kMPS &&
-            self.device().type() != at::kPrivateUse1) {
+            self.device().type() != at::kPrivateUse1 && self.device().type() != at::kXLA) {
           // 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.");
         } else {
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -8936,7 +8936,9 @@ def test_linear_empty(self, device):
         _test_module_empty_input(self, mod, inp)
 
     def test_one_hot(self, device):
-        if self.device_type != 'cuda':  # cuda throws device assert for invalid data
+        # cuda throws device assert for invalid data
+        # xla ignores out of bound indices
+        if self.device_type != 'cuda' and self.device_type != 'xla':
             with self.assertRaises(RuntimeError):
                 torch.nn.functional.one_hot(torch.tensor([3, 4, -1, 0], device=device), -1)
 
