MLBEDSW-7439: Add support for input dims < 4 for ArgMax

- Updated ARG_MAX to support IFM rank less than 4
- Regenerated SUPPORTED_OPS.md

Change-Id: Icd8e72733279413cbea49021325e1ab06fdc6011
Signed-off-by: Johan Alfven <johan.alfven@arm.com>

2 files changed, 3 insertions, 9 deletions

diff --git a/ethosu/vela/tflite_graph_optimiser.py b/ethosu/vela/tflite_graph_optimiser.py
index 518b6db0..e0c7fd2c 100644
--- a/ethosu/vela/tflite_graph_optimiser.py
+++ b/ethosu/vela/tflite_graph_optimiser.py
@@ -41,6 +41,7 @@ from .graph_optimiser_util import needed_total_padding
 from .graph_optimiser_util import set_ifm_ofm_op_shapes
 from .graph_optimiser_util import set_tensor_equivalence
 from .numeric_util import clamp_sigmoid
+from .numeric_util import full_shape
 from .numeric_util import round_away_zero
 from .operation import create_activation_function
 from .operation import ExplicitScaling
@@ -524,7 +525,7 @@ def convert_argmax_to_depthwise_conv_and_max_pool(op, arch, nng):
         op.name = "depthwise_conv_SHL_7"
         op.type = Op.DepthwiseConv2DBias
         op.attrs.update(dw_op_attrs)
-        n, h, w, c = ifm.shape
+        n, h, w, c = full_shape(4, ifm.shape, 1)
         shape = [1, 1, 1, c]
         kernel = np.dstack([2**7] * c)
         op.inputs = []
diff --git a/ethosu/vela/tflite_supported_operators.py b/ethosu/vela/tflite_supported_operators.py
index fd9a9c20..66b9e944 100644
--- a/ethosu/vela/tflite_supported_operators.py
+++ b/ethosu/vela/tflite_supported_operators.py
@@ -316,7 +316,6 @@ class TFLiteSupportedOperators:
         self.specific_constraints[Op.Reshape].append(TFLiteSupportedOperators.constraint_reshape_shape_constant)
 
         # ArgMax specific checks:
-        self.specific_constraints[Op.ArgMax].append(TFLiteSupportedOperators.constraint_argmax_input_dimensions)
         self.specific_constraints[Op.ArgMax].append(TFLiteSupportedOperators.constraint_argmax_axis)
         self.specific_constraints[Op.ArgMax].append(TFLiteSupportedOperators.constraint_argmax_depth)
 
@@ -879,17 +878,11 @@ class TFLiteSupportedOperators:
         inp_dims = len(op.inputs[0].shape)
         axis = op.inputs[1].values
         return (
-            axis in (3, -1),
+            axis in (inp_dims - 1, -1),
             f"Axis is {axis} and number of input dimensions is {inp_dims}",
         )
 
     @staticmethod
-    def constraint_argmax_input_dimensions(op):
-        "Number of input dimensions must be 4"
-        inp_dims = len(op.inputs[0].shape)
-        return inp_dims == 4, f"Number of input dimensions is {inp_dims}"
-
-    @staticmethod
     def constraint_argmax_depth(op):
         "IFM depth must be no greater than 127"
         ifm_depth = op.inputs[0].shape[-1]