[ETHOSN] Improved handling of 5d reshapes (apache#10860)

Resolves an issue with 5d reshapes in the Yolo network and added a
test case. Refactored the reshape tests to use parametrization.

src/relay/backend/contrib/ethosn/ethosn_api.cc

@@ -275,6 +275,11 @@ EthosnError EthosnAPI::Reshape(const Expr& expr, ReshapeParams* params) {
   const auto* input_dtype = reshape->args[0]->checked_type().as<TensorTypeNode>();
   const auto& reshape_attrs = reshape->attrs.as<ReshapeAttrs>();
 
+  if (reshape_attrs->newshape.size() > params->new_shape.size()) {
+    return EthosnError(ErrStrm() << "reshape dimension=" << reshape_attrs->newshape.size()
+                                 << ", reshape dimension must be <= " << params->new_shape.size());
+  }
+
   sl::TensorShape input_tensor_shape = {1, 1, 1, 1};
   sl::DataType input_data_type;
   EthosnError err = Tvm2Npu(input_dtype->shape, &input_tensor_shape);

tests/python/contrib/test_ethosn/test_reshape.py

@@ -36,53 +36,61 @@ def _get_model(input_shape, output_shape, dtype):
 
 @requires_ethosn
 @pytest.mark.parametrize("dtype", ["uint8", "int8"])
-def test_reshape(dtype):
-    trials = [
+@pytest.mark.parametrize(
+    "input_shape, output_shape",
+    [
         ((1, 15, 4, 1), (1, 60)),
         ((1, 15, 4, 1), (1, 30, 2)),
         ((1, 15, 4, 1), (1, 4, 15, 1)),
         ((1, 15, 4, 1), (1, 12, 5, 1)),
         ((1, 15, 4, 1), (1, -1, 2, 1)),
-    ]
-
+    ],
+)
+def test_reshape(dtype, input_shape, output_shape):
     np.random.seed(0)
-    for input_shape, output_shape in trials:
-        inputs = {
-            "a": tvm.nd.array(
-                np.random.randint(
-                    low=np.iinfo(dtype).min,
-                    high=np.iinfo(dtype).max + 1,
-                    size=input_shape,
-                    dtype=dtype,
-                )
+    inputs = {
+        "a": tvm.nd.array(
+            np.random.randint(
+                low=np.iinfo(dtype).min,
+                high=np.iinfo(dtype).max + 1,
+                size=input_shape,
+                dtype=dtype,
             )
-        }
-        outputs = []
-        for npu in [False, True]:
-            model, params = _get_model(input_shape, output_shape, dtype)
-            mod = tei.make_module(model, params)
-            outputs.append(tei.build_and_run(mod, inputs, 1, params, npu=npu))
+        )
+    }
+    outputs = []
+    for npu in [False, True]:
+        model, params = _get_model(input_shape, output_shape, dtype)
+        mod = tei.make_module(model, params)
+        outputs.append(tei.build_and_run(mod, inputs, 1, params, npu=npu))
 
-        tei.verify(outputs, dtype, 1)
+    tei.verify(outputs, dtype, 1)
 
 
 @requires_ethosn
-def test_reshape_failure():
-    trials = [
+@pytest.mark.parametrize(
+    "input_shape, output_shape, dtype, err_msg",
+    [
         (
             (1, 15, 4, 1),
             (1, 15, -2),
             "uint8",
             "reshape dimension=-2, reshape dimension must be >= -1",
         ),
-    ]
-
+        (
+            (1, 1, 4, 1),
+            (1, 1, 2, 2, 1),
+            "uint8",
+            "reshape dimension=5, reshape dimension must be <= 4",
+        ),
+    ],
+)
+def test_reshape_failure(input_shape, output_shape, dtype, err_msg):
     np.random.seed(0)
-    for input_shape, output_shape, dtype, err_msg in trials:
-        model, params = _get_model(input_shape, output_shape, dtype)
-        mod = tei.make_module(model, params)
-        pattern = get_pattern_table("ethos-n")
-        mod = tei.make_module(model, params)
-        mod = relay.transform.MergeComposite(pattern)(mod)
-        mod = tei.make_ethosn_partition(mod["main"].body)
-        tei.test_error(mod, {}, err_msg)
+    model, params = _get_model(input_shape, output_shape, dtype)
+    mod = tei.make_module(model, params)
+    pattern = get_pattern_table("ethos-n")
+    mod = tei.make_module(model, params)
+    mod = relay.transform.MergeComposite(pattern)(mod)
+    mod = tei.make_ethosn_partition(mod["main"].body)
+    tei.test_error(mod, {}, err_msg)