Repeat crash fix (#69)

* Handle empty repeats

* Add repeat to allow list

* Replace numberWithInt with numberWithInteger

Author: abhudev

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

@@ -44,12 +44,21 @@ void set_apparent_shapes(NSMutableArray<NSNumber*> * input_shape,
                          int64_t num_repeat_dims) {
 
 
-  // Set repeats_shape
+  bool repeat_empty = false;
+  if(num_repeat_dims == 0) {
+    num_repeat_dims = num_input_dims;
+    repeat_empty = true;
+  }
 
+  // Set repeats_shape
   repeats_shape = [NSMutableArray<NSNumber*> arrayWithCapacity:num_repeat_dims];
 
-  for(int i = 0; i < num_repeat_dims; i++)
-    repeats_shape[i] = [NSNumber numberWithInt:repeats[i]];
+  for(int i = 0; i < num_repeat_dims; i++) {
+    if(repeat_empty)
+      repeats_shape[i] = [NSNumber numberWithInteger:1];
+    else
+      repeats_shape[i] = [NSNumber numberWithInteger:repeats[i]];
+  }
 
   // If no extension of the shape is needed
   if(num_repeat_dims == num_input_dims) {
@@ -115,7 +124,7 @@ Tensor repeat_mps(const Tensor& self, IntArrayRef repeats) {
                       c10::nullopt);
 
   // Empty output
-  if(zero_tensor)
+  if(zero_tensor || output.numel() == 0)
     return output;
 
   auto stream = at::mps::getCurrentMPSStream();

test/test_mps.py

@@ -6530,6 +6530,12 @@ class TestConsistency(TestCase):
                  'torch.int32',
                  'torch.int64',
                  'torch.uint8'],
+        'repeat': ['torch.float16',
+                   'torch.float32',
+                   'torch.int16',
+                   'torch.int32',
+                   'torch.int64',
+                   'torch.uint8'],
         'repeat_interleave': ['torch.bool',
                               'torch.float16',
                               'torch.float32',
@@ -6699,7 +6705,6 @@ class TestConsistency(TestCase):
         # These were moved from ALLOWLIST to BLOCK as they are not working
         # locally
         'tile': ['torch.float16', 'torch.float32', 'torch.int16', 'torch.int32', 'torch.int64', 'torch.uint8'],
-        'repeat': ['torch.float16', 'torch.float32', 'torch.int16', 'torch.int32', 'torch.int64', 'torch.uint8'],
         '__radd__': ['torch.bool', 'torch.uint8'],
         '__rmul__': ['torch.uint8'],
         'add': ['torch.bool', 'torch.uint8'],