Fix masked_fill broadcasting. (#5688)

test/test_operations.py

@@ -309,6 +309,19 @@ def test_get_xla_tensor(self):
     tx = t.select(1, 12)
     self.assertEqual(tx, sx.data.cpu())
 
+  def test_masked_fill_scalar(self):
+
+    def fn(tensor):
+      # Build a mask from the first line of tensor.
+      # Also, make it have the same rank as the original tensor.
+      mask = tensor[0].ge(0.5).unsqueeze(dim=0)
+      # Call masked_fill.
+      return tensor.masked_fill(mask, 10)
+
+    x = _gen_tensor(2, 2, device=xm.xla_device())
+    x_cpu = x.cpu()
+    self.assertEqual(fn(x_cpu), fn(x))
+
 
 class TestRandom(test_utils.XlaTestCase):
 

torch_xla/csrc/data_ops.cpp

@@ -147,19 +147,19 @@ xla::XlaOp BuildExpand(xla::XlaOp input,
 xla::XlaOp BuildMaskedFillScalar(xla::XlaOp input, xla::XlaOp mask,
                                  xla::XlaOp scalar) {
   const xla::Shape& input_shape = ShapeHelper::ShapeOfXlaOp(input);
-  int64_t input_rank = input_shape.rank();
   const xla::Shape& mask_shape = ShapeHelper::ShapeOfXlaOp(mask);
-  int64_t mask_rank = mask_shape.rank();
-  if (input_rank <= mask_rank) {
-    input = BuildExpand(input, mask_shape.dimensions());
-  } else {
-    mask = BuildExpand(mask, input_shape.dimensions());
+
+  if (!xla::ShapeUtil::Compatible(input_shape, mask_shape)) {
+    xla::Shape shape = XlaHelpers::GetPromotedShape(input_shape, mask_shape);
+    input = BuildExpand(input, shape.dimensions());
+    mask = BuildExpand(mask, shape.dimensions());
   }
+
   xla::XlaOp zero = xla::Zero(mask.builder(), XlaHelpers::TypeOfXlaOp(mask));
   xla::XlaOp mask_pred = xla::Ne(mask, zero);
   xla::XlaOp update_scalar =
       ConvertTo(scalar, ShapeHelper::ShapeOfXlaOp(scalar).element_type(),
-                input_shape.element_type(), nullptr);
+                ShapeHelper::ShapeOfXlaOp(input).element_type(), nullptr);
   return xla::Select(mask_pred, update_scalar, input);
 }