Add data-type promotion to stack. (#7091)

pytorch · May 29, 2024 · 157f8f9 · 157f8f9
1 parent 4e2bbb3
commit 157f8f9
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/test/test_operations.py b/test/test_operations.py
@@ -2005,6 +2005,22 @@ def foo(grad, inp):
 
     self.assertEqual(r, Xr.cpu())
 
+  def test_stack_different_types(self):
+
+    def foo(t0, t1):
+      return torch.stack([t0, t1])
+
+    t0 = torch.rand(10, 10, dtype=torch.bfloat16)
+    t1 = torch.rand(10, 10)
+
+    Xt0 = t0.to(xm.xla_device())
+    Xt1 = t1.to(xm.xla_device())
+
+    r = foo(t0, t1)
+    Xr = foo(Xt0, Xt1)
+
+    self.assertEqual(r, Xr.cpu())
+
 
 class MNISTComparator(nn.Module):
 

diff --git a/torch_xla/csrc/aten_xla_type.cpp b/torch_xla/csrc/aten_xla_type.cpp
@@ -3160,8 +3160,12 @@ at::Tensor XLANativeFunctions::squeeze_copy(const at::Tensor& self,
 
 at::Tensor XLANativeFunctions::stack(at::TensorList tensors, int64_t dim) {
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
+  at::ScalarType result_type = at::native::result_type(tensors);
+  std::vector<at::Tensor> c_tensors(tensors.size());
+  std::transform(tensors.begin(), tensors.end(), c_tensors.begin(),
+                 [=](const at::Tensor& t) { return t.to(result_type); });
   return bridge::AtenFromXlaTensor(
-      tensor_methods::stack(bridge::GetXlaTensors(tensors), dim));
+      tensor_methods::stack(bridge::GetXlaTensors(c_tensors), dim));
 }
 
 at::Tensor XLANativeFunctions::std(const at::Tensor& self, bool unbiased) {