Add support for masked_scatter (#361)

* Add support for masked_scatter

* Fix lintrunner

Author: DenisVieriu97

aten/src/ATen/native/mps/operations/Indexing.mm

@@ -942,6 +942,56 @@ Tensor embedding_dense_backward_mps(
   return masked_fill__mps(self, mask, value.item());
 }
 
+Tensor & masked_scatter__mps(Tensor& self, const Tensor& mask, const Tensor& source) {
+  at::assert_no_internal_overlap(self);
+  TORCH_CHECK(
+      self.scalar_type() == source.scalar_type(),
+      "masked_scatter: expected self and source to have same dtypes but got",
+      self.scalar_type(),
+      " and ",
+      source.scalar_type());
+
+  if (self.numel() == 0) {
+    return self;
+  }
+
+  TORCH_CHECK(mask.scalar_type() == ScalarType::Byte || mask.scalar_type() == ScalarType::Bool,
+              "masked_scatter: expected BoolTensor or ByteTensor for mask");
+
+  auto mask_temp = (mask.dim() == 0)
+    ? c10::MaybeOwned<Tensor>::owned(mask.unsqueeze(0))
+    : c10::MaybeOwned<Tensor>::borrowed(mask);
+  auto self_temp = (self.dim() == 0)
+    ? c10::MaybeOwned<Tensor>::owned(self.unsqueeze(0))
+    : c10::MaybeOwned<Tensor>::borrowed(self);
+
+  // Cannot reassign to mask_temp and self_temp here! if they are
+  // owning and expand_outplace returns a borrow, the returned borrow
+  // would dangle.
+  auto mask_self_expanded = expand_outplace(*mask_temp, *self_temp);
+  auto indices = at::native::expandTensors(
+    *std::get<1>(mask_self_expanded),
+    c10::List<c10::optional<at::Tensor>>({*std::move(std::get<0>(mask_self_expanded))})
+    );
+  // next broadcast all index tensors together
+  try {
+    indices = at::expand_outplace(indices);
+  } catch (std::exception &e) {
+    TORCH_CHECK_INDEX(false, "shape mismatch: indexing tensors could not be broadcast together");
+  }
+
+  if (!indices[0].has_storage() || indices[0].numel() == 0) {
+    return self;
+  }
+
+  return at::index_put_out(
+    self,
+    *std::get<1>(mask_self_expanded),
+    c10::List<c10::optional<at::Tensor>>({*std::move(std::get<0>(mask_self_expanded))}),
+    source.resize_(indices[0].numel())
+  );
+}
+
 REGISTER_DISPATCH(index_stub, &index_kernel_mps);
 REGISTER_DISPATCH(index_put_stub, &index_put_kernel_mps);
 } // namespace at::native

aten/src/ATen/native/native_functions.yaml

@@ -7422,6 +7422,7 @@
   dispatch:
     CPU: masked_scatter__cpu
     CUDA: masked_scatter__cuda
+    MPS: masked_scatter__mps
   autogen: masked_scatter.out
 
 - func: masked_scatter(Tensor self, Tensor mask, Tensor source) -> Tensor

test/test_mps.py

@@ -1003,6 +1003,27 @@ def helper(size, memory_format):
 
         helper((2, 3, 6, 6), torch.contiguous_format)
 
+    def test_masked_scatter(self):
+        def helper(shape):
+            x_mps = torch.randn(shape, device="mps")
+            x_cpu = x_mps.detach().clone().cpu()
+
+            mask_mps = torch.rand(shape, device="mps") < 0.6
+            mask_cpu = mask_mps.detach().clone().cpu()
+
+            y_mps = torch.randn(shape, device="mps")
+            y_cpu = y_mps.detach().clone().cpu()
+
+            y_mps.masked_scatter_(mask_mps, x_mps)
+            y_cpu.masked_scatter_(mask_cpu, x_cpu)
+
+            self.assertEqual(y_mps, y_cpu)
+        helper([2, 5])
+        helper([10, 10])
+        helper([5, 10, 3])
+        helper([10, 5, 10, 3])
+        helper([10, 5, 10, 3, 20])
+
     def test_masked_fill(self):
         device = "mps"
         dtype = torch.float32
@@ -9304,7 +9325,7 @@ class TestConsistency(TestCaseMPS):
         'masked.std': ['f32', 'i16', 'i32', 'i64', 'u8'],
         'masked.var': ['f16', 'f32', 'i16', 'i32', 'i64', 'u8'],
         'masked_fill': ['b8', 'f16', 'f32', 'i16', 'i32', 'i64', 'u8'],
-        'masked_scatter': ['b8', 'f16', 'f32', 'i16', 'i32', 'i64', 'u8'],
+        'masked_scatter': ['i8', 'b8', 'f16', 'f32', 'i16', 'i32', 'i64', 'u8'],
         'masked_select': ['b8', 'f16', 'f32', 'i16', 'i32', 'i64', 'u8'],
         'matmul': ['f32', 'i16', 'i32', 'i64', 'u8'],
         'matrix_exp': ['f32'],
@@ -10425,7 +10446,6 @@ class TestConsistency(TestCaseMPS):
         'lu_unpack': [torch.float32],
         'masked.cumprod': [torch.float32, torch.int16, torch.int32, torch.int64, torch.uint8],
         'masked.median': [torch.float32],
-        'masked_scatter': [torch.bool, torch.float32, torch.float16, torch.int16, torch.int32, torch.int64, torch.uint8],
         'matrix_exp': [torch.float32],
         'mode': [torch.bool, torch.float32, torch.float16, torch.int16, torch.int32, torch.int64, torch.uint8],
         'msort': [torch.bool, torch.float32, torch.float16, torch.int16, torch.int32, torch.int64, torch.uint8],