Fix elementwise arithmetic with zero sized dimensions

Fixes #893.

Author: asuhan

test/cpp/test_aten_xla_tensor.cpp

@@ -189,6 +189,18 @@ TEST_F(AtenXlaTensorTest, TestAddScalarInPlace) {
   });
 }
 
+TEST_F(AtenXlaTensorTest, TestAddZeroSizeDim) {
+  torch::Tensor a = torch::rand({0, 2}, torch::TensorOptions(torch::kFloat));
+  torch::Tensor b = torch::rand({1, 2}, torch::TensorOptions(torch::kFloat));
+  torch::Tensor c = torch::add(a, b);
+  ForEachDevice([&](const torch::Device& device) {
+    torch::Tensor xla_a = CopyToDevice(a, device);
+    torch::Tensor xla_b = CopyToDevice(b, device);
+    torch::Tensor xla_c = torch::add(xla_a, xla_b);
+    AllClose(c, xla_c);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestSub) {
   torch::Tensor a = torch::rand({2, 2}, torch::TensorOptions(torch::kFloat));
   torch::Tensor b = torch::rand({2, 2}, torch::TensorOptions(torch::kFloat));

torch_xla/csrc/helpers.cpp

@@ -307,7 +307,11 @@ xla::Shape XlaHelpers::GetPromotedShape(const xla::Shape& shape1,
     xla::int64 dim2 = shape2_dims[shape2_dims.size() - min_size + i];
     XLA_CHECK(dim1 == dim2 || dim1 == 1 || dim2 == 1)
         << shape1 << " and " << shape2;
-    dimensions.push_back(std::max<xla::int64>(dim1, dim2));
+    if (dim1 == 0 || dim2 == 0) {
+      dimensions.push_back(0);
+    } else {
+      dimensions.push_back(std::max<xla::int64>(dim1, dim2));
+    }
   }
   return xla::ShapeUtil::MakeShape(shape1.element_type(), dimensions);
 }