Fix slice followed by reshape (#237)

Co-authored-by: Denis Vieriu <denisvieriu@mac-3045BF.local>

Author: DenisVieriu97

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

@@ -109,7 +109,8 @@ void binaryOpTensor(const Tensor& self, const Tensor& other, const Scalar& alpha
           newCachedGraph->outputTensor = binaryBlock(newCachedGraph, primaryCastTensor, secondaryCastTensor);
           // Cast output tensor to an expected type if needed, which addresses discrepancy when int64 scalar is added to int32 tensor
           // Output tensor should have been promoted but it remains an int32 tensor
-          if (outputDataType != common_dtype) {
+          if (outputDataType != common_dtype ||
+             [newCachedGraph->outputTensor dataType] != getMPSDataType(outputDataType)) {
             newCachedGraph->outputTensor = castMPSTensor(mpsGraph, newCachedGraph->outputTensor, outputDataType);
           }
         }

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

@@ -536,11 +536,26 @@ bool canSliceViewTensor(const Tensor& src, MPSShape *mpsShape) {
   std::vector<int64_t> src_view_shape = getViewShape(src, mpsShape);
   int src_ndim_view = src_view_shape.size();
   if (src_ndim_base == src_ndim_view) {
-    for (const auto i : c10::irange(src_ndim_base)) {
+    for (const auto i: c10::irange(src_ndim_base)) {
       if (src_view_shape[i] > src_base_shape[i]) {
         return false;
       }
     }
+  } else {
+    // Detect slice followed by reshape cases, e.g (1,4800,2) -> (1,4800)
+    bool allDimsEqual = true;
+    auto min_ndim = std::min(src_ndim_base, src_ndim_view);
+    for (const auto i: c10::irange(min_ndim)) {
+      if (src_view_shape[i] > src_base_shape[i]) {
+        return false;
+      }
+      else if (src_view_shape[i] != src_base_shape[i]) {
+        allDimsEqual = false;
+      }
+    }
+    if (allDimsEqual) {
+      return false;
+    }
   }
   return true;
 }

test/test_mps.py

@@ -1617,9 +1617,23 @@ def test_slice_reshape(self):
 
         x = x[:,3:].view(2, 3, 4, 1)
         x_cpu = x_cpu[:,3:].view(2, 3, 4, 1)
+        self.assertEqual(x, x_cpu)
 
+        x = x + 2
+        x_cpu = x_cpu + 2
         self.assertEqual(x, x_cpu)
 
+    def test_slice_reshape_contg_view(self):
+        import torch
+
+        x_mps = torch.randn(1, 4800, 2, device="mps")
+        x_cpu = x_mps.detach().clone().cpu()
+
+        r_mps = x_mps + 2
+        r_cpu = x_cpu + 2
+
+        self.assertEqual(r_mps, r_cpu)
+
     def test_view_slice(self):
         # https://github.com/pytorch/pytorch/issues/83995
         NUM_SAMPLES = 60