Cortex-M backend: Verify output shape before rewriting AdaptiveAvgPool (pytorch#19935)

The pass only does a naive rewrite, so check that the output shape
actually matches after the rewrite before doing it.

Signed-off-by: Erik Lundell <erik.lundell@arm.com>

Author: Erik-Lundell

backends/cortex_m/passes/decompose_mean_pass.py

@@ -25,11 +25,21 @@ def call_operator(
         meta: NodeMetadata,
     ) -> ProxyValue:
         if op == torch.ops.aten.adaptive_avg_pool2d.default:
-            op = torch.ops.aten.avg_pool2d.default
-            input_tensor = cast(torch.Tensor, args[0])
-            shape = input_tensor.data.shape
+            input_tensor = cast(ProxyValue, args[0]).to_tensor()
+            shape = input_tensor.shape
             stride = [1, 1]
             kernel_size = [shape[-2], shape[-1]]
-            args = (args[0], kernel_size, stride, [0, 0], 0, 0)
 
+            new_args = (args[0], kernel_size, stride, [0, 0], 0, 0)
+
+            adaptive_output = torch.ops.aten.adaptive_avg_pool2d.default(
+                input_tensor, *args[1:]
+            )
+            avg_pool_output = torch.ops.aten.avg_pool2d.default(
+                input_tensor, *new_args[1:]
+            )
+
+            if adaptive_output.shape == avg_pool_output.shape:
+                new_op = torch.ops.aten.avg_pool2d.default
+                return super().call_operator(new_op, new_args, kwargs, meta)
         return super().call_operator(op, args, kwargs, meta)