pytorch
diff --git a/‎aten/src/ATen/native/mps/operations/Indexing.mm‎
Lines changed: 22 additions & 1 deletion b/‎aten/src/ATen/native/mps/operations/Indexing.mm‎
Lines changed: 22 additions & 1 deletion
diff --git a/‎torch/testing/_internal/common_mps.py‎
Lines changed: 1 addition & 1 deletion b/‎torch/testing/_internal/common_mps.py‎
Lines changed: 1 addition & 1 deletion
@@ -512,7 +512,28 @@ Tensor flip_mps(const Tensor& self, IntArrayRef dims) {
     return;
   }
 
-  TORCH_CHECK(source.scalar_type() != ScalarType::Long, "index_add(): Expected non int64 dtype for source.");
+  bool use_deterministic_algorithm = globalContext().deterministicAlgorithms();
+
+  // TODO: Do not use deterministic algorithm for long/complex but rather implement it as Metal shader
+  use_deterministic_algorithm |= source.scalar_type() == ScalarType::Long;
+  use_deterministic_algorithm |= c10::isComplexType(source.scalar_type());
+
+  if (use_deterministic_algorithm) {
+    if (!result.is_same(self)) {
+      result.copy_(self);
+    }
+    torch::List<std::optional<Tensor>> indices;
+    indices.reserve(dim + 1);
+    for (const auto i : c10::irange(dim)) {
+      indices.emplace_back();
+    }
+    indices.emplace_back(index.to(at::kLong));
+    const Tensor result_ = (result.dim() == 0) ? result.view(1) : result;
+    const Tensor source_ = (source.dim() == 0) ? source.view(1) : source;
+    result_.index_put_(indices, source_.mul(alpha), true);
+    return;
+  }
+
   auto casted_type = isFloatingType(source.scalar_type()) ? ScalarType::Float : ScalarType::Int;
 
   struct CachedGraph : public MPSCachedGraph {
 
@@ -74,6 +74,7 @@ def mps_ops_modifier(
             "H",
             "hsplit",
             "imag",
+            "index_add",
             "index_copy",
             "index_select",
             "index_put",
@@ -419,7 +420,6 @@ def mps_ops_modifier(
             ],
             # Unsupported dtypes
             "histc": [torch.float16, torch.bfloat16],
-            "index_add": [torch.int64],
             # GEMM on MPS is not supported for integral types
             "nn.functional.linear": [
                 torch.int16,