Added zero check to inverse op, resolving crash seen in inverse & matrix_pow tests (#236)

Co-authored-by: Siddharth Kotapati <sidk@Siddharths-MacBook-Pro.local>

Author: 2 people

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

@@ -24,6 +24,10 @@
     MPSStream* stream = getCurrentMPSStream();
     info.zero_();
 
+    if (A.numel() == 0) {
+        return;
+    }
+
     struct CachedGraph : public MPSCachedGraph
     {
         CachedGraph(MPSGraph *graph) : MPSCachedGraph(graph) {}

test/test_mps.py

@@ -9402,8 +9402,6 @@ class TestConsistency(TestCase):
         'nonzero': [torch.bool, torch.uint8, torch.float16],
         'median': [torch.float32, torch.int16, torch.int32, torch.uint8, torch.int16],
         'sgn': [torch.bool],
-        'linalg.inv': [torch.float32],
-        'linalg.inv_ex': [torch.float32],
         'linalg.matrix_power': [torch.float32],
         'nn.functional.interpolate': [torch.float32],
         'resize_': [torch.bool, torch.float16, torch.float32, torch.int16, torch.int32, torch.int64, torch.uint8],