Native channel shuffle floating point exception (pytorch#144010)

Fixes pytorch#142453

Added TORCH_CHECKS to prevent the user from using the native_channel_shuffle function incorrectly and getting a "Floating point exception (core dumped)"

Pull Request resolved: pytorch#144010
Approved by: https://github.com/albanD

Author: abcarlisle

aten/src/ATen/native/ChanelShuffle.cpp

@@ -20,6 +20,17 @@
 namespace at::native {
 
 Tensor channel_shuffle_cpu(const Tensor& self, int64_t groups) {
+  TORCH_CHECK(self.dim() > 2,
+              "channel_shuffle expects input with > 2 dims, but got input with sizes ",
+              self.sizes());
+  int64_t c = self.size(1);
+  TORCH_CHECK(groups > 0,
+              "Number of groups to divide channels in must be positive.",
+              " Value of groups:", groups);
+  TORCH_CHECK((c % groups) == 0,
+              "Number of channels must be divisible by groups. Got ",
+              c, " channels and ", groups, " groups.");
+
   Tensor output;
   if (self.numel() == 0) {
     output = self.alias();

test/test_nn.py

@@ -6400,6 +6400,24 @@ def test_channel_shuffle_return_alias_of_self(self):
         output = torch.nn.ChannelShuffle(groups)(input_tensor)
         torch.testing.assert_close(output, input_tensor)
 
+    def test_channel_shuffle_input_checks(self):
+        input_tensor = torch.rand([1, 3, 2, 2])
+        with self.assertRaisesRegex(RuntimeError,
+                                    "Number of groups to divide channels in must be positive.*"):
+            groups = 0
+            torch.native_channel_shuffle(input_tensor, groups)
+
+        with self.assertRaisesRegex(RuntimeError,
+                                    "Number of channels must be divisible by groups.*"):
+            groups = 2
+            torch.native_channel_shuffle(input_tensor, groups)
+
+        with self.assertRaisesRegex(RuntimeError,
+                                    "channel_shuffle expects input with > 2 dims,.*"):
+            input_tensor = torch.rand([1, 2])
+            groups = 2
+            torch.native_channel_shuffle(input_tensor, groups)
+
     @skipIfTorchDynamo("TorchDynamo fails here for unknown reasons")
     def test_native_channel_shuffle_return_alias_of_self(self):
         groups = 3